Feat: support Gemini models

.pre-commit-config.yaml

@@ -55,7 +55,7 @@ repos:
             verbose: true
         -   id: type-checking-hook
             name: Run type checking (mypy)
-            entry: bash -xc './.venv-wayflowcore/bin/mypy -- "$@"' --
+            entry: bash -xc 'mypy --config ./pyproject.toml -- "$@"' --
             language: system
             types_or: [python]
             # We exclude the fuzz tests from this check

docs/wayflowcore/source/core/api/llmmodels.rst

@@ -91,6 +91,16 @@ OpenAI Models
 .. _openaimodel:
 .. autoclass:: wayflowcore.models.openaimodel.OpenAIModel
 
+Gemini Models
+~~~~~~~~~~~~~
+
+.. _geminimodel:
+.. autoclass:: wayflowcore.models.geminimodel.GeminiModel
+
+.. autoclass:: wayflowcore.models.geminimodel.GeminiApiKeyAuth
+
+.. autoclass:: wayflowcore.models.geminimodel.GeminiCloudAuth
+
 Ollama Models
 ~~~~~~~~~~~~~
 

docs/wayflowcore/source/core/changelog.rst

@@ -26,6 +26,13 @@ New features
   Added support for converting `MessageSummarizationTransform` and `ConversationSummarizationTransform` between Agent Spec and Wayflow. Similarly for Datastores (`OracleDatabaseDatastore`, `PostgreSQLDatabaseDatastore`).
   You can now declare your agents with summarization transforms and summary caching in Agent Spec and run them in WayFlow.
 
+* **Gemini models (Vertex AI + AI Studio):**
+
+  Added ``GeminiModel`` to run Gemini models via Google Vertex AI and Google AI Studio.
+  For more information read the :doc:`API Reference on LLM models <api/llmmodels>` and the guide on
+  :doc:`how to use LLMs from different providers <howtoguides/llm_from_different_providers>`.
+
+
 Improvements
 ^^^^^^^^^^^^
 

docs/wayflowcore/source/core/code_examples/example_initialize_llms.py

@@ -130,6 +130,57 @@
         ca_file="/etc/certs/ca.pem",
     )
 # .. openai-compatible-tls-end
+# .. gemini-start
+from wayflowcore.models import GeminiApiKeyAuth, GeminiCloudAuth, GeminiModel
+
+if __name__ == "__main__":
+
+    generation_config = LlmGenerationConfig(max_tokens=256, temperature=0.7, top_p=0.9)
+
+    llm = GeminiModel(
+        model_id="gemini-2.5-flash",
+        auth=GeminiApiKeyAuth(),
+        generation_config=generation_config,
+    )
+
+    vertex_llm = GeminiModel(
+        model_id="gemini-2.0-flash-lite",
+        auth=GeminiCloudAuth(
+            # Often still required even when ADC supplies the credentials.
+            project_id="my-gcp-project",
+            location="global",
+            # Optional when Google Application Default Credentials are already configured.
+            vertex_credentials="path/to/service-account.json",
+        ),
+        generation_config=generation_config,
+    )
+# .. gemini-end
+
+# .. gemini-llmfactory-start
+if __name__ == "__main__":
+    GEMINI_AISTUDIO_CONFIG = {
+        "model_type": "gemini",
+        "model_id": "gemini-2.5-flash",
+        "auth": {"type": "api_key"},
+        "generation_config": {"max_tokens": 256, "temperature": 0.7, "top_p": 0.9},
+    }
+    GEMINI_VERTEX_CONFIG = {
+        "model_type": "gemini",
+        "model_id": "gemini-2.0-flash-lite",
+        "auth": {
+            "type": "cloud",
+            # Often still required even when ADC supplies the credentials.
+            "project_id": "my-gcp-project",
+            "location": "global",
+            # Optional when Google Application Default Credentials are already configured.
+            "vertex_credentials": "path/to/service-account.json",
+        },
+        "generation_config": {"max_tokens": 256, "temperature": 0.7, "top_p": 0.9},
+    }
+
+    llm = LlmModelFactory.from_config(GEMINI_AISTUDIO_CONFIG)
+    vertex_llm = LlmModelFactory.from_config(GEMINI_VERTEX_CONFIG)
+# .. gemini-llmfactory-end
 
 # .. ollama-start
 from wayflowcore.models import OllamaModel
@@ -159,7 +210,15 @@
 # .. ollama-llmfactory-end
 
 # .. recap:
-from wayflowcore.models import OCIGenAIModel, OllamaModel, OpenAIModel, VllmModel
+from wayflowcore.models import (
+    GeminiApiKeyAuth,
+    GeminiCloudAuth,
+    GeminiModel,
+    OCIGenAIModel,
+    OllamaModel,
+    OpenAIModel,
+    VllmModel,
+)
 from wayflowcore.models.llmgenerationconfig import LlmGenerationConfig
 
 OCIGENAI_MODEL_ID = "cohere.command-r-plus"
@@ -202,6 +261,24 @@
         generation_config=generation_config,
     )
 
+    llm = GeminiModel(
+        model_id="gemini-2.5-flash",
+        auth=GeminiApiKeyAuth(),
+        generation_config=generation_config,
+    )
+
+    vertex_llm = GeminiModel(
+        model_id="gemini-2.0-flash-lite",
+        auth=GeminiCloudAuth(
+            # Often still required even when ADC supplies the credentials.
+            project_id="my-gcp-project",
+            location="global",
+            # Optional when Google Application Default Credentials are already configured.
+            vertex_credentials="path/to/service-account.json",
+        ),
+        generation_config=generation_config,
+    )
+
     OLLAMA_MODEL_ID = "llama2-7b"
     OLLAMA_HOST_PORT = "localhost:11434"  # default is 11434 if omitted
 

docs/wayflowcore/source/core/howtoguides/llm_from_different_providers.rst

@@ -5,6 +5,7 @@ How to Use LLMs from Different LLM Providers
 WayFlow supports several LLM API providers. The available LLMs are:
 
 - :ref:`OpenAIModel <openaimodel>`
+- :ref:`GeminiModel <geminimodel>`
 - :ref:`OpenAICompatibleModel <openaicompatiblemodel>`
 - :ref:`OCIGenAIModel <ocigenaimodel>`
 - :ref:`VllmModel <vllmmodel>`
@@ -221,6 +222,53 @@ OpenAI Model is powered by `OpenAI <https://platform.openai.com/docs/models>`_.
         :start-after: .. openai-llmfactory-start
         :end-before: .. openai-llmfactory-end
 
+Gemini Model
+------------
+
+`Gemini <https://deepmind.google/technologies/gemini/>`_ models can be accessed through
+Google AI Studio or Google Vertex AI.
+
+**Parameters**
+
+.. option:: model_id: str
+
+  Gemini model name as exposed by the selected service. For example, use the
+  model name listed in Google AI Studio when authenticating with
+  ``GeminiApiKeyAuth``, and use the model name listed in Vertex AI when
+  authenticating with ``GeminiCloudAuth``.
+
+.. option:: auth: GeminiApiKeyAuth | GeminiCloudAuth
+
+  Required authentication configuration for Gemini. Use ``GeminiApiKeyAuth()`` if
+  you want the runtime to read ``GEMINI_API_KEY`` from the environment, or
+  ``GeminiCloudAuth(...)`` for Vertex AI.
+
+.. option:: generation_config: dict, optional
+
+  Default parameters for text generation with this model.
+
+**Examples**
+
+.. literalinclude:: ../code_examples/example_initialize_llms.py
+    :language: python
+    :start-after: .. gemini-start
+    :end-before: .. gemini-end
+
+.. collapse:: Equivalent code example utilizing the LlmModelFactory class.
+
+    .. literalinclude:: ../code_examples/example_initialize_llms.py
+        :language: python
+        :start-after: .. gemini-llmfactory-start
+        :end-before: .. gemini-llmfactory-end
+
+**Notes**
+
+- Google AI Studio auth may read ``GEMINI_API_KEY`` from the environment.
+- The important part is the service-side model name. AI Studio and Vertex AI may expose different Gemini model names or availability, so use the exact model name shown by the service you are connecting to.
+- Vertex AI auth accepts an optional service-account JSON file path or an inline service-account JSON object.
+- When ``vertex_credentials`` is omitted, LiteLLM/Google auth may rely on Application Default Credentials (ADC), such as ``GOOGLE_APPLICATION_CREDENTIALS``, the local ADC file created by ``gcloud auth application-default login``, or an attached service account.
+- ADC does not always provide a default Vertex project, so ``project_id`` should usually still be set explicitly unless your environment already resolves it.
+
 vLLM Model
 ----------
 

wayflowcore/constraints/constraints.txt

@@ -11,3 +11,4 @@ mcp==1.24.0
 pyagentspec==26.1.0
 opentelemetry-api==1.36.0
 opentelemetry-sdk==1.36.0
+litellm==1.81.13

wayflowcore/constraints/constraints_dev.txt

@@ -11,3 +11,4 @@ mcp==1.24.0
 pyagentspec==26.2.0.dev3 # Main branch of pyagentspec
 opentelemetry-api==1.36.0
 opentelemetry-sdk==1.36.0
+litellm==1.81.13

wayflowcore/setup.cfg

@@ -54,6 +54,9 @@ filterwarnings =
     ignore:SSL verification is disabled. This is not recommended for production environments.:UserWarning
     ignore:Found no serialization plugin to serialize the object of type:UserWarning
     ignore:Found no deserialization plugin to deserialize the object of type:UserWarning
+    # Litellm always passes enable_cleanup_closed=True to aiohttp; aiohttp emits
+    # this deprecation warning on Python versions where CPython fixed the SSL leak bug.
+    ignore:enable_cleanup_closed ignored because:DeprecationWarning
     # The following ignores are needed to run tests on python 3.14
     # This deprecation warning is thrown by the oci package (even in latest today's version 2.164.0)
     ignore:'asyncio\.iscoroutinefunction' is deprecated and slated for removal in Python 3\.16:DeprecationWarning

wayflowcore/setup.py

@@ -1,4 +1,4 @@
-# Copyright © 2025 Oracle and/or its affiliates.
+# Copyright © 2025, 2026 Oracle and/or its affiliates.
 #
 # This software is under the Apache License 2.0
 # (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
@@ -69,6 +69,7 @@ def read(file_name):
         "exceptiongroup>=1.0.2",
         "uvicorn>=0.23.1",
         "fastapi>=0.116.2,<1.0.0",
+        "litellm>=1.81.13,<1.82.0",
         # 4rth party dependencies version bounds, for CVE patching
         "annotated-types>=0.6.0",
         "certifi>=2025.4.26",

wayflowcore/src/wayflowcore/models/__init__.py

@@ -9,6 +9,7 @@
     TaggedMessageChunkType,
     TaggedMessageChunkTypeWithTokenUsage,
 )
+from .geminimodel import GeminiApiKeyAuth, GeminiCloudAuth, GeminiModel
 from .llmgenerationconfig import LlmGenerationConfig
 from .llmmodel import LlmCompletion, LlmModel, Prompt
 from .llmmodelfactory import LlmModelFactory
@@ -53,4 +54,7 @@
     "ServingMode",
     "ModelProvider",
     "OciAPIType",
+    "GeminiApiKeyAuth",
+    "GeminiCloudAuth",
+    "GeminiModel",
 ]

wayflowcore/src/wayflowcore/models/_openaihelpers/_chatcompletions_processor.py

@@ -1,4 +1,4 @@
-# Copyright © 2025 Oracle and/or its affiliates.
+# Copyright © 2025, 2026 Oracle and/or its affiliates.
 #
 # This software is under the Apache License 2.0
 # (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0) or Universal Permissive License
@@ -66,7 +66,7 @@ def _convert_message_into_openai_message_dict(
                 raise ValueError(
                     "Invalid tool request. A tool request message should only contain text contents"
                 )
-            converted_message = {
+            converted_message: Dict[str, Any] = {
                 "role": "assistant",
                 "tool_calls": [
                     {
@@ -84,6 +84,8 @@ def _convert_message_into_openai_message_dict(
             }
             if m.content:
                 converted_message["content"] = m.content
+            if m._extra_content is not None:
+                converted_message["extra_content"] = m._extra_content
             return [converted_message]
         elif m.tool_result:
             if len(m.contents):
@@ -112,7 +114,13 @@ def _convert_message_into_openai_message_dict(
                 else:
                     raise RuntimeError(f"Unsupported content type: {content.__class__.__name__}")
 
-            return [{"role": role, "content": all_contents if len(all_contents) else ""}]
+            assistant_message: Dict[str, Any] = {
+                "role": role,
+                "content": all_contents if len(all_contents) else "",
+            }
+            if role == "assistant" and m._extra_content is not None:
+                assistant_message["extra_content"] = m._extra_content
+            return [assistant_message]
 
     def _convert_prompt(self, prompt: "Prompt", supports_tool_role: bool) -> Dict[str, Any]:
         payload_arguments: Dict[str, Any] = {
@@ -169,6 +177,7 @@ def _convert_openai_response_into_message(self, response: Any) -> "Message":
                     for tc in extracted_message["tool_calls"]
                 ],
                 role="assistant",
+                _extra_content=extracted_message.get("extra_content"),
             )
         else:
             # content might be empty when certain models (like gemini) decide