feat(integrations): add OCI Generative AI provider
Adds OCIGenAILlm under integrations/oci/, for Google Gemini and other models hosted on Oracle Cloud Infrastructure Generative AI. Optional install: pip install google-adk[oci]. LLMRegistry auto-routing and the google.adk.models import surface are preserved. The OpenAI-compatible transport from the source PR (OCIGenAIOpenAILlm) is not taken. It reimplemented the message, tool and response conversion plus the streaming loop that OpenAILlm already provides; the right form is a small subclass overriding the OpenAI client, which cannot live in integrations/ while OpenAILlm is still experimental. It can land separately once that settles. The OCI client is now built once per instance rather than per request, so a call no longer re-reads the OCI config from disk. Merge https://github.com/google/adk-python/pull/5285 Closes #5069 Co-authored-by: George Weale <gweale@google.com> COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5285 from fede-kamel:feat/oci-generative-ai 0230acc0a93b7e43014f2ef3a8b89de463a50bd8 PiperOrigin-RevId: 955453382
This commit is contained in:
Fede Kamelhar
committed by
Copybara-Service
Copybara-Service
parent
94832a5151
commit
625ef1aa69
@@ -196,6 +196,9 @@ optional-dependencies.mcp = [
|
||||
"anyio>=4.9,<5",
|
||||
"mcp>=1.24,<2",
|
||||
]
|
||||
optional-dependencies.oci = [
|
||||
"oci>=2.126", # OCI Generative AI native SDK (OCIGenAILlm)
|
||||
]
|
||||
optional-dependencies.otel-gcp = [
|
||||
"opentelemetry-instrumentation-google-genai>=0.7b1,<1",
|
||||
"opentelemetry-instrumentation-grpc>=0.43b0,<1",
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""OCI Generative AI integration.
|
||||
|
||||
Model providers for Google Gemini and other models hosted on Oracle Cloud
|
||||
Infrastructure (OCI) Generative AI. Install with: pip install google-adk[oci]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import typing
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ._oci_genai_llm import OCIGenAILlm
|
||||
|
||||
_lazy_imports = {
|
||||
"OCIGenAILlm": "._oci_genai_llm",
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> typing.Any:
|
||||
if name in _lazy_imports:
|
||||
import importlib
|
||||
|
||||
module = importlib.import_module(_lazy_imports[name], __name__)
|
||||
return getattr(module, name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return list(_lazy_imports.keys())
|
||||
@@ -0,0 +1,653 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""OCI Generative AI integration for ADK models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from functools import cached_property
|
||||
import importlib.util
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
from typing import AsyncGenerator
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from google.genai import types
|
||||
from typing_extensions import override
|
||||
|
||||
if not TYPE_CHECKING and importlib.util.find_spec("oci") is None:
|
||||
raise ImportError(
|
||||
"OCI Generative AI support requires: pip install google-adk[oci]"
|
||||
"\nOr: pip install oci"
|
||||
)
|
||||
|
||||
from google.adk.models.base_llm import BaseLlm
|
||||
from google.adk.models.llm_response import LlmResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
|
||||
__all__ = ["OCIGenAILlm"]
|
||||
|
||||
logger = logging.getLogger("google_adk." + __name__)
|
||||
|
||||
|
||||
def _to_oci_role(role: Optional[str]) -> str:
|
||||
"""Map ADK content role to OCI GenAI role string."""
|
||||
if role in ("model", "assistant"):
|
||||
return "ASSISTANT"
|
||||
return "USER"
|
||||
|
||||
|
||||
def _build_response_format(
|
||||
cfg: types.GenerateContentConfig, oci_models: Any
|
||||
) -> Optional[Any]:
|
||||
"""Map google.genai response config to OCI ResponseFormat.
|
||||
|
||||
- ``response_schema`` (Pydantic class, dict, or genai Schema) →
|
||||
``JsonSchemaResponseFormat`` (strict structured output).
|
||||
- ``response_mime_type == "application/json"`` only →
|
||||
``JsonObjectResponseFormat``.
|
||||
- ``response_mime_type == "text/plain"`` → ``TextResponseFormat`` (default
|
||||
behaviour; only emitted when explicitly requested).
|
||||
"""
|
||||
schema = cfg.response_schema
|
||||
mime = cfg.response_mime_type or ""
|
||||
|
||||
if schema is not None:
|
||||
schema_dict: dict[str, Any]
|
||||
if hasattr(schema, "model_json_schema"):
|
||||
# Pydantic v2 model class
|
||||
schema_dict = schema.model_json_schema()
|
||||
elif hasattr(schema, "to_json_dict"):
|
||||
# google.genai Schema instance
|
||||
schema_dict = schema.to_json_dict()
|
||||
elif isinstance(schema, dict):
|
||||
schema_dict = schema
|
||||
else:
|
||||
return None
|
||||
return oci_models.JsonSchemaResponseFormat(
|
||||
type="JSON_SCHEMA",
|
||||
json_schema=oci_models.ResponseJsonSchema(
|
||||
name=schema_dict.get("title", "response"),
|
||||
description=schema_dict.get("description"),
|
||||
schema=schema_dict,
|
||||
is_strict=True,
|
||||
),
|
||||
)
|
||||
|
||||
if mime == "application/json":
|
||||
return oci_models.JsonObjectResponseFormat(type="JSON_OBJECT")
|
||||
if mime == "text/plain":
|
||||
return oci_models.TextResponseFormat(type="TEXT")
|
||||
return None
|
||||
|
||||
|
||||
def _media_blocks_for_part(part: types.Part) -> list[Any]:
|
||||
"""Map a multimodal Part (inline_data / file_data) to OCI ChatContent blocks.
|
||||
|
||||
OCI Generative AI Inference (/20231130/) accepts ImageContent / AudioContent /
|
||||
VideoContent / DocumentContent, each carrying a URL in ``{kind}_url.url``.
|
||||
Inline bytes are wrapped as ``data:<mime>;base64,<...>``; file_data passes
|
||||
the ``file_uri`` through.
|
||||
|
||||
Returns an empty list for parts that have no media payload.
|
||||
"""
|
||||
import oci.generative_ai_inference.models as oci_models
|
||||
|
||||
url: Optional[str] = None
|
||||
mime: Optional[str] = None
|
||||
|
||||
if part.inline_data and part.inline_data.data is not None:
|
||||
mime = part.inline_data.mime_type or "application/octet-stream"
|
||||
raw = part.inline_data.data
|
||||
if isinstance(raw, (bytes, bytearray)):
|
||||
encoded = base64.b64encode(bytes(raw)).decode("ascii")
|
||||
else:
|
||||
encoded = str(raw)
|
||||
url = f"data:{mime};base64,{encoded}"
|
||||
elif part.file_data and part.file_data.file_uri:
|
||||
url = part.file_data.file_uri
|
||||
mime = part.file_data.mime_type
|
||||
|
||||
if not url:
|
||||
return []
|
||||
|
||||
category = (mime or "").split("/", 1)[0].lower()
|
||||
if category == "image":
|
||||
return [
|
||||
oci_models.ImageContent(
|
||||
type="IMAGE", image_url=oci_models.ImageUrl(url=url)
|
||||
)
|
||||
]
|
||||
if category == "audio":
|
||||
return [
|
||||
oci_models.AudioContent(
|
||||
type="AUDIO", audio_url=oci_models.AudioUrl(url=url)
|
||||
)
|
||||
]
|
||||
if category == "video":
|
||||
return [
|
||||
oci_models.VideoContent(
|
||||
type="VIDEO", video_url=oci_models.VideoUrl(url=url)
|
||||
)
|
||||
]
|
||||
# Documents (application/pdf, text/*, etc.) and any other mime
|
||||
return [
|
||||
oci_models.DocumentContent(
|
||||
type="DOCUMENT", document_url=oci_models.DocumentUrl(url=url)
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _content_to_oci_message(content: types.Content) -> Any:
|
||||
"""Convert an ADK Content object to an OCI GenAI message.
|
||||
|
||||
OCI GenAI uses:
|
||||
- ``UserMessage`` for user turns
|
||||
- ``AssistantMessage`` for model turns (may include ``FunctionCall`` items
|
||||
in ``tool_calls``)
|
||||
- ``ToolMessage`` for tool results (function_response parts)
|
||||
"""
|
||||
import oci.generative_ai_inference.models as oci_models
|
||||
|
||||
text_parts: list[str] = []
|
||||
media_blocks: list[Any] = []
|
||||
tool_calls: list[Any] = []
|
||||
tool_results: list[tuple[str, str]] = [] # (tool_call_id, result_text)
|
||||
|
||||
for part in content.parts or []:
|
||||
if part.text:
|
||||
text_parts.append(part.text)
|
||||
elif part.function_call:
|
||||
# FunctionCall is the OCI subtype of ToolCall that carries name+arguments
|
||||
tool_calls.append(
|
||||
oci_models.FunctionCall(
|
||||
id=part.function_call.id or "",
|
||||
type=oci_models.FunctionCall.TYPE_FUNCTION,
|
||||
name=part.function_call.name,
|
||||
arguments=json.dumps(part.function_call.args or {}),
|
||||
)
|
||||
)
|
||||
elif part.function_response:
|
||||
result = part.function_response.response or {}
|
||||
tool_results.append((
|
||||
part.function_response.id or "",
|
||||
json.dumps(result) if isinstance(result, dict) else str(result),
|
||||
))
|
||||
elif part.inline_data or part.file_data:
|
||||
media_blocks.extend(_media_blocks_for_part(part))
|
||||
|
||||
role = _to_oci_role(content.role)
|
||||
|
||||
# Tool results map to ToolMessage (one per result)
|
||||
if tool_results:
|
||||
call_id, result_text = tool_results[0]
|
||||
return oci_models.ToolMessage(
|
||||
role=oci_models.ToolMessage.ROLE_TOOL,
|
||||
tool_call_id=call_id,
|
||||
content=[oci_models.TextContent(type="TEXT", text=result_text)],
|
||||
)
|
||||
|
||||
if role == "ASSISTANT":
|
||||
oci_content: list[Any] = []
|
||||
if text_parts:
|
||||
oci_content.append(
|
||||
oci_models.TextContent(type="TEXT", text="\n".join(text_parts))
|
||||
)
|
||||
return oci_models.AssistantMessage(
|
||||
role=oci_models.AssistantMessage.ROLE_ASSISTANT,
|
||||
content=oci_content,
|
||||
tool_calls=tool_calls or None,
|
||||
)
|
||||
|
||||
user_content: list[Any] = []
|
||||
if text_parts:
|
||||
user_content.append(
|
||||
oci_models.TextContent(type="TEXT", text="\n".join(text_parts))
|
||||
)
|
||||
user_content.extend(media_blocks)
|
||||
return oci_models.UserMessage(
|
||||
role=oci_models.UserMessage.ROLE_USER,
|
||||
content=user_content,
|
||||
)
|
||||
|
||||
|
||||
def _oci_response_to_llm_response(response: Any) -> LlmResponse:
|
||||
"""Convert an OCI GenAI chat response to an LlmResponse."""
|
||||
chat_response = response.data.chat_response
|
||||
parts: list[types.Part] = []
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
reasoning_tokens = 0
|
||||
|
||||
if hasattr(chat_response, "usage"):
|
||||
usage = chat_response.usage
|
||||
input_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(usage, "completion_tokens", 0) or 0
|
||||
details = getattr(usage, "completion_tokens_details", None)
|
||||
if details is not None:
|
||||
reasoning_tokens = getattr(details, "reasoning_tokens", 0) or 0
|
||||
|
||||
if hasattr(chat_response, "choices") and chat_response.choices:
|
||||
choice = chat_response.choices[0]
|
||||
message = getattr(choice, "message", None)
|
||||
if message:
|
||||
# Text content
|
||||
for block in getattr(message, "content", None) or []:
|
||||
if hasattr(block, "text") and block.text:
|
||||
parts.append(types.Part.from_text(text=block.text))
|
||||
|
||||
# Tool calls — OCI returns FunctionCall objects directly in tool_calls
|
||||
for fc in getattr(message, "tool_calls", None) or []:
|
||||
args: dict[str, Any] = {}
|
||||
try:
|
||||
args = json.loads(fc.arguments) if fc.arguments else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
args = {}
|
||||
part = types.Part.from_function_call(
|
||||
name=fc.name,
|
||||
args=args,
|
||||
)
|
||||
if part.function_call is not None:
|
||||
part.function_call.id = getattr(fc, "id", "") or ""
|
||||
parts.append(part)
|
||||
|
||||
return LlmResponse(
|
||||
content=types.Content(role="model", parts=parts),
|
||||
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=input_tokens,
|
||||
candidates_token_count=output_tokens,
|
||||
total_token_count=input_tokens + output_tokens,
|
||||
thoughts_token_count=reasoning_tokens or None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _function_declaration_to_oci_tool(
|
||||
fn: types.FunctionDeclaration,
|
||||
) -> Any:
|
||||
"""Convert an ADK FunctionDeclaration to an OCI GenAI Tool."""
|
||||
import oci.generative_ai_inference.models as oci_models
|
||||
|
||||
parameters: dict[str, Any] = {"type": "object", "properties": {}}
|
||||
if fn.parameters_json_schema:
|
||||
parameters = fn.parameters_json_schema
|
||||
elif fn.parameters and fn.parameters.properties:
|
||||
props = {}
|
||||
for k, v in fn.parameters.properties.items():
|
||||
props[k] = v.model_dump(by_alias=True, exclude_none=True)
|
||||
parameters = {
|
||||
"type": "object",
|
||||
"properties": props,
|
||||
}
|
||||
if fn.parameters.required:
|
||||
parameters["required"] = fn.parameters.required
|
||||
|
||||
return oci_models.FunctionDefinition(
|
||||
type=oci_models.FunctionDefinition.TYPE_FUNCTION,
|
||||
name=fn.name,
|
||||
description=fn.description or "",
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
|
||||
class OCIGenAILlm(BaseLlm):
|
||||
"""Integration with OCI Generative AI models.
|
||||
|
||||
Supports models hosted on Oracle Cloud Infrastructure Generative AI service,
|
||||
including Meta Llama, Google Gemini, Google Gemma, and other GenericChat
|
||||
compatible models.
|
||||
|
||||
Example usage::
|
||||
|
||||
from google.adk.integrations.oci import OCIGenAILlm
|
||||
from google.adk.agents import LlmAgent
|
||||
|
||||
agent = LlmAgent(
|
||||
model=OCIGenAILlm(
|
||||
model="google.gemini-2.0-flash-001",
|
||||
compartment_id="ocid1.compartment.oc1...",
|
||||
),
|
||||
...
|
||||
)
|
||||
|
||||
Attributes:
|
||||
model: OCI model ID (e.g. ``google.gemini-2.0-flash-001``). Used as the
|
||||
``model_id`` for on-demand serving. For dedicated serving, set
|
||||
``endpoint_id`` instead; ``model`` is then informational only.
|
||||
endpoint_id: Dedicated endpoint OCID (``ocid1.generativeaiendpoint...``).
|
||||
When set, requests use ``DedicatedServingMode``; otherwise on-demand
|
||||
mode is used. Falls back to ``OCI_ENDPOINT_ID`` env var when not set.
|
||||
compartment_id: OCI compartment OCID. Falls back to the
|
||||
``OCI_COMPARTMENT_ID`` environment variable when not set.
|
||||
service_endpoint: OCI Generative AI service endpoint URL. Defaults to
|
||||
the us-chicago-1 endpoint or ``OCI_SERVICE_ENDPOINT`` env var.
|
||||
auth_type: OCI authentication type. One of ``API_KEY`` (default),
|
||||
``INSTANCE_PRINCIPAL``, or ``RESOURCE_PRINCIPAL``.
|
||||
auth_profile: Config profile to use for ``API_KEY`` auth (default:
|
||||
``DEFAULT``).
|
||||
auth_file_location: Path to the OCI config file used for ``API_KEY``
|
||||
auth (default: ``~/.oci/config``).
|
||||
max_tokens: Maximum number of tokens to generate (default: 2048).
|
||||
reasoning_effort: Reasoning-token budget for reasoning-capable models.
|
||||
One of ``"NONE"``, ``"MINIMAL"``, ``"LOW"``, ``"MEDIUM"``, ``"HIGH"``,
|
||||
or ``None`` (default — let OCI pick). Honoured by GPT-5 family,
|
||||
Gemini 2.5, Grok reasoning variants, and Cohere Command-A-Reasoning;
|
||||
ignored by non-reasoning models. The single most impactful cost knob
|
||||
for reasoning models — ``"LOW"`` typically cuts reasoning-token spend
|
||||
5-10× vs the default.
|
||||
"""
|
||||
|
||||
model: str = "google.gemini-2.5-flash"
|
||||
endpoint_id: Optional[str] = None
|
||||
compartment_id: Optional[str] = None
|
||||
service_endpoint: Optional[str] = None
|
||||
auth_type: str = "API_KEY"
|
||||
auth_profile: str = "DEFAULT"
|
||||
auth_file_location: str = "~/.oci/config"
|
||||
max_tokens: int = 2048
|
||||
reasoning_effort: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
@override
|
||||
def supported_models(cls) -> list[str]:
|
||||
return [
|
||||
r"meta\.llama-.*",
|
||||
r"google\.gemini-.*",
|
||||
r"google\.gemma-.*",
|
||||
r"xai\.grok-.*",
|
||||
r"mistralai\.mistral-.*",
|
||||
r"mistralai\.mixtral-.*",
|
||||
r"nvidia\..*",
|
||||
]
|
||||
|
||||
@override
|
||||
async def generate_content_async(
|
||||
self,
|
||||
llm_request: LlmRequest,
|
||||
stream: bool = False,
|
||||
) -> AsyncGenerator[LlmResponse, None]:
|
||||
if stream:
|
||||
async for response in self._generate_content_streaming(llm_request):
|
||||
yield response
|
||||
else:
|
||||
response = await asyncio.to_thread(self._call_oci, llm_request)
|
||||
yield _oci_response_to_llm_response(response)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _resolve_compartment_id(self) -> str:
|
||||
compartment_id = self.compartment_id or os.environ.get("OCI_COMPARTMENT_ID")
|
||||
if not compartment_id:
|
||||
raise ValueError(
|
||||
"compartment_id must be set on OCIGenAILlm or via the"
|
||||
" OCI_COMPARTMENT_ID environment variable."
|
||||
)
|
||||
return compartment_id
|
||||
|
||||
def _resolve_service_endpoint(self) -> str:
|
||||
return (
|
||||
self.service_endpoint
|
||||
or os.environ.get("OCI_SERVICE_ENDPOINT")
|
||||
or "https://inference.generativeai.us-chicago-1.oci.oraclecloud.com"
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def _oci_client(self) -> Any:
|
||||
return self._build_client(self._resolve_service_endpoint())
|
||||
|
||||
def _build_client(self, service_endpoint: str) -> Any:
|
||||
"""Create an OCI GenerativeAiInferenceClient from auth config."""
|
||||
import oci
|
||||
import oci.auth.signers
|
||||
import oci.generative_ai_inference
|
||||
|
||||
if self.auth_type == "INSTANCE_PRINCIPAL":
|
||||
signer = oci.auth.signers.InstancePrincipalsSecurityTokenSigner()
|
||||
return oci.generative_ai_inference.GenerativeAiInferenceClient(
|
||||
config={},
|
||||
signer=signer,
|
||||
service_endpoint=service_endpoint,
|
||||
)
|
||||
elif self.auth_type == "RESOURCE_PRINCIPAL":
|
||||
signer = oci.auth.signers.get_resource_principals_signer()
|
||||
return oci.generative_ai_inference.GenerativeAiInferenceClient(
|
||||
config={},
|
||||
signer=signer,
|
||||
service_endpoint=service_endpoint,
|
||||
)
|
||||
else: # API_KEY (default)
|
||||
config = oci.config.from_file(
|
||||
file_location=self.auth_file_location,
|
||||
profile_name=self.auth_profile,
|
||||
)
|
||||
return oci.generative_ai_inference.GenerativeAiInferenceClient(
|
||||
config=config,
|
||||
service_endpoint=service_endpoint,
|
||||
)
|
||||
|
||||
def _build_chat_details(
|
||||
self, llm_request: LlmRequest, is_stream: bool = False
|
||||
) -> Any:
|
||||
"""Build OCI ChatDetails from an LlmRequest."""
|
||||
import oci.generative_ai_inference.models as oci_models
|
||||
|
||||
messages = [_content_to_oci_message(c) for c in llm_request.contents or []]
|
||||
|
||||
# Prepend SystemMessage when a system instruction is present
|
||||
if llm_request.config and llm_request.config.system_instruction:
|
||||
si = llm_request.config.system_instruction
|
||||
if isinstance(si, str) and si:
|
||||
messages = [
|
||||
oci_models.SystemMessage(
|
||||
role=oci_models.SystemMessage.ROLE_SYSTEM,
|
||||
content=[oci_models.TextContent(type="TEXT", text=si)],
|
||||
)
|
||||
] + messages
|
||||
|
||||
# Convert tool declarations if present
|
||||
oci_tools: Optional[list[Any]] = None
|
||||
if llm_request.config and llm_request.config.tools:
|
||||
first_tool = llm_request.config.tools[0]
|
||||
if (
|
||||
isinstance(first_tool, types.Tool)
|
||||
and first_tool.function_declarations
|
||||
):
|
||||
oci_tools = [
|
||||
_function_declaration_to_oci_tool(fn)
|
||||
for fn in first_tool.function_declarations
|
||||
]
|
||||
|
||||
chat_request_kwargs: dict[str, Any] = dict(
|
||||
api_format=oci_models.BaseChatRequest.API_FORMAT_GENERIC,
|
||||
messages=messages,
|
||||
max_tokens=self.max_tokens,
|
||||
)
|
||||
|
||||
# Sampling and decoding parameters from llm_request.config.
|
||||
cfg = getattr(llm_request, "config", None)
|
||||
if cfg is not None:
|
||||
if cfg.max_output_tokens is not None:
|
||||
chat_request_kwargs["max_tokens"] = cfg.max_output_tokens
|
||||
if cfg.temperature is not None:
|
||||
chat_request_kwargs["temperature"] = cfg.temperature
|
||||
if cfg.top_p is not None:
|
||||
chat_request_kwargs["top_p"] = cfg.top_p
|
||||
if cfg.top_k is not None:
|
||||
chat_request_kwargs["top_k"] = int(cfg.top_k)
|
||||
if cfg.frequency_penalty is not None:
|
||||
chat_request_kwargs["frequency_penalty"] = cfg.frequency_penalty
|
||||
if cfg.presence_penalty is not None:
|
||||
chat_request_kwargs["presence_penalty"] = cfg.presence_penalty
|
||||
if cfg.seed is not None:
|
||||
chat_request_kwargs["seed"] = cfg.seed
|
||||
if cfg.stop_sequences:
|
||||
chat_request_kwargs["stop"] = list(cfg.stop_sequences)
|
||||
|
||||
# Structured-output: response_schema (Pydantic / dict / google.genai
|
||||
# Schema) → JsonSchemaResponseFormat. response_mime_type alone (without
|
||||
# a schema) → JsonObjectResponseFormat (free-form JSON).
|
||||
response_format = _build_response_format(cfg, oci_models)
|
||||
if response_format is not None:
|
||||
chat_request_kwargs["response_format"] = response_format
|
||||
|
||||
# Constructor-level reasoning_effort applies regardless of per-request cfg.
|
||||
if self.reasoning_effort is not None:
|
||||
chat_request_kwargs["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
if oci_tools:
|
||||
chat_request_kwargs["tools"] = oci_tools
|
||||
if is_stream:
|
||||
chat_request_kwargs["is_stream"] = True
|
||||
chat_request_kwargs["stream_options"] = oci_models.StreamOptions(
|
||||
is_include_usage=True
|
||||
)
|
||||
|
||||
return oci_models.ChatDetails(
|
||||
compartment_id=self._resolve_compartment_id(),
|
||||
serving_mode=self._build_serving_mode(oci_models),
|
||||
chat_request=oci_models.GenericChatRequest(**chat_request_kwargs),
|
||||
)
|
||||
|
||||
def _build_serving_mode(self, oci_models: Any) -> Any:
|
||||
endpoint_id = self.endpoint_id or os.environ.get("OCI_ENDPOINT_ID")
|
||||
if endpoint_id:
|
||||
return oci_models.DedicatedServingMode(endpoint_id=endpoint_id)
|
||||
return oci_models.OnDemandServingMode(model_id=self.model)
|
||||
|
||||
def _call_oci(self, llm_request: LlmRequest) -> Any:
|
||||
"""Synchronous non-streaming OCI GenAI call, run in a thread pool."""
|
||||
chat_details = self._build_chat_details(llm_request, is_stream=False)
|
||||
logger.debug("Sending request to OCI GenAI: model=%s", self.model)
|
||||
return self._oci_client.chat(chat_details)
|
||||
|
||||
def _call_oci_stream(self, llm_request: LlmRequest) -> list[dict[str, Any]]:
|
||||
"""Synchronous streaming call — collects all SSE event dicts in a thread.
|
||||
|
||||
The OCI SDK wraps an SSE response in ``oci._vendor.sseclient.SSEClient``
|
||||
when ``is_stream=True`` is set on the request body. Each event's
|
||||
``data`` field is an OpenAI-compatible JSON chunk or the sentinel
|
||||
``[DONE]``.
|
||||
"""
|
||||
chat_details = self._build_chat_details(llm_request, is_stream=True)
|
||||
logger.debug("Sending streaming request to OCI GenAI: model=%s", self.model)
|
||||
response = self._oci_client.chat(chat_details)
|
||||
|
||||
chunks: list[dict[str, Any]] = []
|
||||
try:
|
||||
for event in response.data.events():
|
||||
raw = getattr(event, "data", None)
|
||||
if not raw or raw.strip() == "[DONE]":
|
||||
break
|
||||
try:
|
||||
chunks.append(json.loads(raw))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
logger.debug("Could not parse SSE event data: %r", raw)
|
||||
finally:
|
||||
close = getattr(response.data, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
return chunks
|
||||
|
||||
async def _generate_content_streaming(
|
||||
self, llm_request: LlmRequest
|
||||
) -> AsyncGenerator[LlmResponse, None]:
|
||||
"""Yield partial then final LlmResponse from an OCI SSE stream.
|
||||
|
||||
The OCI SDK is synchronous, so every SSE chunk is collected in a background
|
||||
thread before any response is emitted. Partial responses are therefore not
|
||||
delivered incrementally.
|
||||
"""
|
||||
chunks = await asyncio.to_thread(self._call_oci_stream, llm_request)
|
||||
|
||||
text_acc: str = ""
|
||||
tool_acc: dict[int, dict[str, Any]] = {}
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
reasoning_tokens: int = 0
|
||||
|
||||
for chunk in chunks:
|
||||
# Usage chunk (camelCase per OCI GenAI /20231130/ schema).
|
||||
usage = chunk.get("usage")
|
||||
if usage:
|
||||
input_tokens = usage.get("promptTokens", 0) or 0
|
||||
output_tokens = usage.get("completionTokens", 0) or 0
|
||||
details = usage.get("completionTokensDetails") or {}
|
||||
reasoning_tokens = details.get("reasoningTokens", 0) or 0
|
||||
continue
|
||||
|
||||
message = chunk.get("message")
|
||||
if not message:
|
||||
continue
|
||||
|
||||
# Text content: list of {type: TEXT, text: ...} blocks.
|
||||
for block in message.get("content") or []:
|
||||
if block.get("type") == "TEXT" and block.get("text"):
|
||||
delta_text = block["text"]
|
||||
text_acc += delta_text
|
||||
yield LlmResponse(
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[types.Part.from_text(text=delta_text)],
|
||||
),
|
||||
partial=True,
|
||||
)
|
||||
|
||||
# Tool calls: OCI emits the whole call in one chunk for Gemini, but
|
||||
# accumulate name/arguments defensively in case other providers split
|
||||
# them across events.
|
||||
for tc_idx, tc in enumerate(message.get("toolCalls") or []):
|
||||
idx = tc.get("index", tc_idx)
|
||||
if idx not in tool_acc:
|
||||
tool_acc[idx] = {"id": "", "name": "", "arguments": ""}
|
||||
if tc.get("id"):
|
||||
tool_acc[idx]["id"] = tc["id"]
|
||||
if tc.get("name"):
|
||||
tool_acc[idx]["name"] = tc["name"]
|
||||
if tc.get("arguments"):
|
||||
tool_acc[idx]["arguments"] += tc["arguments"]
|
||||
|
||||
# Build final aggregated response
|
||||
all_parts: list[types.Part] = []
|
||||
if text_acc:
|
||||
all_parts.append(types.Part.from_text(text=text_acc))
|
||||
for tc in sorted(tool_acc.values(), key=lambda x: x.get("name", "")):
|
||||
args: dict[str, Any] = {}
|
||||
try:
|
||||
args = json.loads(tc["arguments"]) if tc["arguments"] else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
args = {}
|
||||
part = types.Part.from_function_call(name=tc["name"], args=args)
|
||||
if part.function_call is not None:
|
||||
part.function_call.id = tc["id"]
|
||||
all_parts.append(part)
|
||||
|
||||
yield LlmResponse(
|
||||
content=types.Content(role="model", parts=all_parts),
|
||||
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=input_tokens,
|
||||
candidates_token_count=output_tokens,
|
||||
total_token_count=input_tokens + output_tokens,
|
||||
thoughts_token_count=reasoning_tokens or None,
|
||||
),
|
||||
partial=False,
|
||||
)
|
||||
@@ -26,6 +26,7 @@ from .llm_response import LlmResponse
|
||||
from .registry import LLMRegistry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from google.adk.integrations.oci._oci_genai_llm import OCIGenAILlm
|
||||
from google.adk.labs.openai import OpenAILlm
|
||||
|
||||
from .anthropic_llm import AnthropicGenerateContentConfig
|
||||
@@ -91,6 +92,18 @@ _LAZY_PROVIDERS: dict[str, tuple[list[str], str]] = {
|
||||
],
|
||||
'lite_llm',
|
||||
),
|
||||
'OCIGenAILlm': (
|
||||
[
|
||||
r'meta\.llama-.*',
|
||||
r'google\.gemini-.*',
|
||||
r'google\.gemma-.*',
|
||||
r'xai\.grok-.*',
|
||||
r'mistralai\.mistral-.*',
|
||||
r'mistralai\.mixtral-.*',
|
||||
r'nvidia\..*',
|
||||
],
|
||||
'google.adk.integrations.oci._oci_genai_llm',
|
||||
),
|
||||
}
|
||||
|
||||
for _name, (_patterns, _module) in _LAZY_PROVIDERS.items():
|
||||
|
||||
@@ -0,0 +1,681 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Integration tests for OCIGenAILlm against live OCI Generative AI service.
|
||||
|
||||
Required environment variables:
|
||||
OCI_COMPARTMENT_ID — OCI compartment OCID
|
||||
OCI_REGION — OCI region (default: us-chicago-1)
|
||||
|
||||
Optional:
|
||||
OCI_AUTH_TYPE — API_KEY | INSTANCE_PRINCIPAL | RESOURCE_PRINCIPAL
|
||||
(default: API_KEY)
|
||||
OCI_AUTH_PROFILE — OCI config profile (default: DEFAULT)
|
||||
OCI_AUTH_FILE — path to OCI config file (default: ~/.oci/config)
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
from google.adk.integrations.oci._oci_genai_llm import OCIGenAILlm
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Skip the entire module when required env vars are absent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# OCI tests do not use any Google backend (GOOGLE_AI / Vertex AI).
|
||||
# Override the autouse llm_backend fixture from the integration conftest so
|
||||
# these tests are not duplicated across backends.
|
||||
@pytest.fixture(autouse=True)
|
||||
def llm_backend():
|
||||
yield
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.environ.get("OCI_COMPARTMENT_ID"),
|
||||
reason=(
|
||||
"OCI integration tests require OCI_COMPARTMENT_ID to be set. "
|
||||
"Set OCI_COMPARTMENT_ID (and optionally OCI_REGION) to run."
|
||||
),
|
||||
)
|
||||
|
||||
_COMPARTMENT_ID = os.environ.get("OCI_COMPARTMENT_ID", "")
|
||||
_REGION = os.environ.get("OCI_REGION", "us-chicago-1")
|
||||
_SERVICE_ENDPOINT = (
|
||||
f"https://inference.generativeai.{_REGION}.oci.oraclecloud.com"
|
||||
)
|
||||
_AUTH_TYPE = os.environ.get("OCI_AUTH_TYPE", "API_KEY")
|
||||
_AUTH_PROFILE = os.environ.get("OCI_AUTH_PROFILE", "DEFAULT")
|
||||
_AUTH_FILE = os.environ.get("OCI_AUTH_FILE", "~/.oci/config")
|
||||
|
||||
_GEMINI_MODEL = "google.gemini-2.5-flash"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gemini_llm() -> OCIGenAILlm:
|
||||
return OCIGenAILlm(
|
||||
model=_GEMINI_MODEL,
|
||||
compartment_id=_COMPARTMENT_ID,
|
||||
service_endpoint=_SERVICE_ENDPOINT,
|
||||
auth_type=_AUTH_TYPE,
|
||||
auth_profile=_AUTH_PROFILE,
|
||||
auth_file_location=_AUTH_FILE,
|
||||
max_tokens=512,
|
||||
)
|
||||
|
||||
|
||||
def _simple_request(
|
||||
model: str, text: str = "Reply with one word: hello."
|
||||
) -> LlmRequest:
|
||||
return LlmRequest(
|
||||
model=model,
|
||||
contents=[Content(role="user", parts=[Part.from_text(text=text)])],
|
||||
)
|
||||
|
||||
|
||||
def _request_with_system(model: str) -> LlmRequest:
|
||||
return LlmRequest(
|
||||
model=model,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="What is your name?")],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
system_instruction=(
|
||||
"Your name is Oracle. Always introduce yourself as Oracle."
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _request_with_tool(model: str) -> LlmRequest:
|
||||
return LlmRequest(
|
||||
model=model,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="What is the weather in Chicago?")],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
tools=[
|
||||
types.Tool(
|
||||
function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name="get_weather",
|
||||
description="Get the current weather for a city.",
|
||||
parameters=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
properties={
|
||||
"city": types.Schema(
|
||||
type=types.Type.STRING,
|
||||
description="The city name.",
|
||||
)
|
||||
},
|
||||
required=["city"],
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gemini (google.gemini-2.0-flash-001) tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_text(gemini_llm):
|
||||
"""Gemini on OCI returns a non-empty text response."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_simple_request(_GEMINI_MODEL), stream=False
|
||||
)
|
||||
]
|
||||
assert len(responses) == 1
|
||||
assert responses[0].content.role == "model"
|
||||
assert responses[0].content.parts
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_usage_metadata(gemini_llm):
|
||||
"""Response includes token usage metadata."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_simple_request(_GEMINI_MODEL), stream=False
|
||||
)
|
||||
]
|
||||
usage = responses[0].usage_metadata
|
||||
assert usage.prompt_token_count > 0
|
||||
assert usage.candidates_token_count > 0
|
||||
assert usage.total_token_count == (
|
||||
usage.prompt_token_count + usage.candidates_token_count
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_with_system_instruction(gemini_llm):
|
||||
"""System instruction is respected."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_request_with_system(_GEMINI_MODEL), stream=False
|
||||
)
|
||||
]
|
||||
text = responses[0].content.parts[0].text.lower()
|
||||
assert "oracle" in text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_tool_call(gemini_llm):
|
||||
"""Gemini returns a function call when a tool is provided."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_request_with_tool(_GEMINI_MODEL), stream=False
|
||||
)
|
||||
]
|
||||
parts = responses[0].content.parts
|
||||
function_calls = [p for p in parts if p.function_call]
|
||||
assert function_calls, "Expected at least one function call in the response"
|
||||
fc = function_calls[0].function_call
|
||||
assert fc.name == "get_weather"
|
||||
assert "city" in fc.args
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_streaming_text(gemini_llm):
|
||||
"""Streaming returns partial chunks followed by a final non-partial response."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_simple_request(_GEMINI_MODEL), stream=True
|
||||
)
|
||||
]
|
||||
assert responses, "Expected at least one response chunk"
|
||||
partial_responses = [r for r in responses if r.partial]
|
||||
final_responses = [r for r in responses if not r.partial]
|
||||
assert partial_responses, "Expected at least one partial (streaming) chunk"
|
||||
assert (
|
||||
len(final_responses) == 1
|
||||
), "Expected exactly one final (non-partial) response"
|
||||
full_text = "".join(
|
||||
p.text for r in partial_responses for p in r.content.parts or [] if p.text
|
||||
)
|
||||
assert full_text.strip(), "Streamed text should be non-empty"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_streaming_usage_metadata(gemini_llm):
|
||||
"""Final streaming response includes token usage metadata."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_simple_request(_GEMINI_MODEL), stream=True
|
||||
)
|
||||
]
|
||||
final = next(r for r in responses if not r.partial)
|
||||
usage = final.usage_metadata
|
||||
assert usage is not None
|
||||
assert usage.prompt_token_count > 0
|
||||
assert usage.candidates_token_count > 0
|
||||
assert usage.total_token_count == (
|
||||
usage.prompt_token_count + usage.candidates_token_count
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_streaming_tool_call(gemini_llm):
|
||||
"""Streaming returns a function call when a tool is provided."""
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_request_with_tool(_GEMINI_MODEL), stream=True
|
||||
)
|
||||
]
|
||||
final = next(r for r in responses if not r.partial)
|
||||
parts = final.content.parts or []
|
||||
function_calls = [p for p in parts if p.function_call]
|
||||
assert (
|
||||
function_calls
|
||||
), "Expected at least one function call in the streaming response"
|
||||
fc = function_calls[0].function_call
|
||||
assert fc.name == "get_weather"
|
||||
assert "city" in fc.args
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_generate_content_concurrent(gemini_llm):
|
||||
"""Multiple concurrent non-streaming requests complete independently."""
|
||||
import asyncio
|
||||
|
||||
async def single_call(text: str) -> str:
|
||||
responses = [
|
||||
r
|
||||
async for r in gemini_llm.generate_content_async(
|
||||
_simple_request(_GEMINI_MODEL, text=text), stream=False
|
||||
)
|
||||
]
|
||||
return responses[0].content.parts[0].text
|
||||
|
||||
results = await asyncio.gather(
|
||||
*[single_call(f"Reply with the number {i} only.") for i in range(3)]
|
||||
)
|
||||
assert len(results) == 3
|
||||
for result in results:
|
||||
assert result.strip(), "Each concurrent response should be non-empty"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_multi_turn(gemini_llm):
|
||||
"""Multi-turn conversation passes history correctly."""
|
||||
history = [
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="My favourite colour is blue.")],
|
||||
),
|
||||
Content(
|
||||
role="model",
|
||||
parts=[Part.from_text(text="Got it, blue is a great colour!")],
|
||||
),
|
||||
]
|
||||
follow_up = Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="What is my favourite colour?")],
|
||||
)
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=history + [follow_up],
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
text = responses[0].content.parts[0].text.lower()
|
||||
assert "blue" in text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-provider on-demand smoke tests
|
||||
#
|
||||
# Skipped unless the corresponding model env var is set so cost stays opt-in.
|
||||
# Set OCI_LLAMA_MODEL / OCI_MISTRAL_MODEL / OCI_GROK_MODEL / OCI_NVIDIA_MODEL
|
||||
# to a model id available in your tenancy/region (e.g. "meta.llama-3.3-70b-instruct").
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _provider_llm(env_var: str) -> "OCIGenAILlm | None":
|
||||
model_id = os.environ.get(env_var)
|
||||
if not model_id:
|
||||
return None
|
||||
return OCIGenAILlm(
|
||||
model=model_id,
|
||||
compartment_id=_COMPARTMENT_ID,
|
||||
service_endpoint=_SERVICE_ENDPOINT,
|
||||
auth_type=_AUTH_TYPE,
|
||||
auth_profile=_AUTH_PROFILE,
|
||||
auth_file_location=_AUTH_FILE,
|
||||
max_tokens=256,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OCI_LLAMA_MODEL"),
|
||||
reason="Set OCI_LLAMA_MODEL=<meta.llama-...-id> to enable.",
|
||||
)
|
||||
async def test_llama_on_demand_generate_text():
|
||||
llm = _provider_llm("OCI_LLAMA_MODEL")
|
||||
responses = [
|
||||
r
|
||||
async for r in llm.generate_content_async(
|
||||
_simple_request(llm.model), stream=False
|
||||
)
|
||||
]
|
||||
assert len(responses) == 1
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OCI_MISTRAL_MODEL"),
|
||||
reason="Set OCI_MISTRAL_MODEL=<mistralai...-id> to enable.",
|
||||
)
|
||||
async def test_mistral_on_demand_generate_text():
|
||||
llm = _provider_llm("OCI_MISTRAL_MODEL")
|
||||
responses = [
|
||||
r
|
||||
async for r in llm.generate_content_async(
|
||||
_simple_request(llm.model), stream=False
|
||||
)
|
||||
]
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OCI_GROK_MODEL"),
|
||||
reason="Set OCI_GROK_MODEL=<xai.grok-...-id> to enable.",
|
||||
)
|
||||
async def test_grok_on_demand_generate_text():
|
||||
llm = _provider_llm("OCI_GROK_MODEL")
|
||||
responses = [
|
||||
r
|
||||
async for r in llm.generate_content_async(
|
||||
_simple_request(llm.model), stream=False
|
||||
)
|
||||
]
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OCI_NVIDIA_MODEL"),
|
||||
reason="Set OCI_NVIDIA_MODEL=<nvidia...-id> to enable.",
|
||||
)
|
||||
async def test_nvidia_on_demand_generate_text():
|
||||
llm = _provider_llm("OCI_NVIDIA_MODEL")
|
||||
responses = [
|
||||
r
|
||||
async for r in llm.generate_content_async(
|
||||
_simple_request(llm.model), stream=False
|
||||
)
|
||||
]
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dedicated serving mode
|
||||
#
|
||||
# Set OCI_DEDICATED_ENDPOINT_ID=ocid1.generativeaiendpoint.oc1... to enable.
|
||||
# OCI_DEDICATED_MODEL is informational; defaults to the dedicated endpoint's
|
||||
# bound model (the SDK ignores `model` when serving_mode is dedicated).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_DEDICATED_ENDPOINT_ID = os.environ.get("OCI_DEDICATED_ENDPOINT_ID", "")
|
||||
_DEDICATED_MODEL = os.environ.get(
|
||||
"OCI_DEDICATED_MODEL", "meta.llama-3.3-70b-instruct"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dedicated_llm() -> OCIGenAILlm:
|
||||
return OCIGenAILlm(
|
||||
model=_DEDICATED_MODEL,
|
||||
endpoint_id=_DEDICATED_ENDPOINT_ID,
|
||||
compartment_id=_COMPARTMENT_ID,
|
||||
service_endpoint=_SERVICE_ENDPOINT,
|
||||
auth_type=_AUTH_TYPE,
|
||||
auth_profile=_AUTH_PROFILE,
|
||||
auth_file_location=_AUTH_FILE,
|
||||
max_tokens=256,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not _DEDICATED_ENDPOINT_ID,
|
||||
reason=(
|
||||
"Set OCI_DEDICATED_ENDPOINT_ID to a dedicated endpoint OCID to enable."
|
||||
),
|
||||
)
|
||||
async def test_dedicated_generate_content_text(dedicated_llm):
|
||||
responses = [
|
||||
r
|
||||
async for r in dedicated_llm.generate_content_async(
|
||||
_simple_request(_DEDICATED_MODEL), stream=False
|
||||
)
|
||||
]
|
||||
assert len(responses) == 1
|
||||
assert responses[0].content.parts[0].text.strip()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
not _DEDICATED_ENDPOINT_ID,
|
||||
reason=(
|
||||
"Set OCI_DEDICATED_ENDPOINT_ID to a dedicated endpoint OCID to enable."
|
||||
),
|
||||
)
|
||||
async def test_dedicated_generate_content_streaming(dedicated_llm):
|
||||
chunks = []
|
||||
async for r in dedicated_llm.generate_content_async(
|
||||
_simple_request(_DEDICATED_MODEL, text="Count from 1 to 3."),
|
||||
stream=True,
|
||||
):
|
||||
chunks.append(r)
|
||||
assert len(chunks) >= 2 # at least one partial + one final
|
||||
final = chunks[-1]
|
||||
assert final.usage_metadata is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sampling parameters (live)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_max_output_tokens_caps_response(gemini_llm):
|
||||
"""max_output_tokens is honoured: completion tokens never exceed the budget.
|
||||
|
||||
Note: Gemini 2.5 spends part of the budget on reasoning tokens before any
|
||||
visible output. We pick a budget large enough to leave some text but small
|
||||
enough to clearly cap an alphabet-recitation response, and we assert on the
|
||||
reported token count rather than character count (which is flaky).
|
||||
"""
|
||||
budget = 64
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(
|
||||
text="Recite the alphabet, A through Z, comma separated."
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(max_output_tokens=budget),
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
um = responses[0].usage_metadata
|
||||
assert um.candidates_token_count is not None
|
||||
assert um.candidates_token_count <= budget
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_low_temperature_deterministic_with_seed(gemini_llm):
|
||||
"""temperature=0 + seed should yield consistent answers across two calls."""
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="Reply with exactly: 'green'")],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(temperature=0.0, seed=12345),
|
||||
)
|
||||
call_a = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
call_b = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
assert "green" in call_a[0].content.parts[0].text.lower()
|
||||
assert "green" in call_b[0].content.parts[0].text.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_stop_sequences_terminate_output(gemini_llm):
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="Print: APPLE | BANANA | CHERRY")],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.0, stop_sequences=["BANANA"]
|
||||
),
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
text = responses[0].content.parts[0].text
|
||||
assert "BANANA" not in text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multimodal: inline image (live)
|
||||
#
|
||||
# Uses a tiny 1x1 red PNG so the request is cheap. Gemini 2.5 Flash on OCI
|
||||
# supports image inputs via ImageContent.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_red_png_1x1() -> bytes:
|
||||
"""Generate a guaranteed-valid 1x1 red PNG with correct CRCs."""
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
sig = b"\x89PNG\r\n\x1a\n"
|
||||
|
||||
def chunk(t: bytes, d: bytes) -> bytes:
|
||||
return (
|
||||
struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d))
|
||||
)
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) # 1x1 RGB
|
||||
idat = zlib.compress(b"\x00\xff\x00\x00") # filter byte + RGB(255,0,0)
|
||||
return sig + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")
|
||||
|
||||
|
||||
_TINY_RED_PNG = _make_red_png_1x1()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_inline_image_input(gemini_llm):
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(
|
||||
text=(
|
||||
"What is the dominant colour of this image? "
|
||||
"Reply with just the colour name."
|
||||
)
|
||||
),
|
||||
Part(
|
||||
inline_data=types.Blob(
|
||||
mime_type="image/png", data=_TINY_RED_PNG
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.0, max_output_tokens=256
|
||||
),
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
parts = responses[0].content.parts
|
||||
assert parts, "Expected the model to produce a visible answer"
|
||||
text = parts[0].text.lower()
|
||||
assert "red" in text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Structured output: response_schema (live)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_response_schema_returns_valid_json(gemini_llm):
|
||||
schema = {
|
||||
"title": "CityFact",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"country": {"type": "string"},
|
||||
},
|
||||
"required": ["city", "country"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[Part.from_text(text="Give me a fact about Paris.")],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
response_mime_type="application/json",
|
||||
response_schema=schema,
|
||||
temperature=0.0,
|
||||
),
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
raw = responses[0].content.parts[0].text
|
||||
payload = json.loads(raw)
|
||||
assert "city" in payload
|
||||
assert "country" in payload
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reasoning-token surfacing (live)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_reasoning_tokens_reported(gemini_llm):
|
||||
"""Gemini 2.5 emits reasoningTokens in completionTokensDetails — surface them."""
|
||||
request = LlmRequest(
|
||||
model=_GEMINI_MODEL,
|
||||
contents=[
|
||||
Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_text(
|
||||
text=(
|
||||
"If a train travels 60km in 30 minutes, what is its"
|
||||
" speed?"
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
config=types.GenerateContentConfig(temperature=0.0),
|
||||
)
|
||||
responses = [r async for r in gemini_llm.generate_content_async(request)]
|
||||
um = responses[0].usage_metadata
|
||||
assert um is not None
|
||||
# Reasoning tokens are optional; assert it's an int when present
|
||||
assert um.thoughts_token_count is None or um.thoughts_token_count > 0
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user