diff --git a/openkyrozen/providers/google.py b/openkyrozen/providers/google.py index 7f95949..67dcaa9 100644 --- a/openkyrozen/providers/google.py +++ b/openkyrozen/providers/google.py @@ -1,140 +1,382 @@ + from __future__ import annotations -import openkyrozen.providers.usage as usage_ledger +import json import os import sys import time -from typing import Any, Iterator +from collections.abc import Iterator, Mapping +from typing import Any + +import openkyrozen.providers.usage as usage_ledger from openkyrozen.providers.base import LLMProvider -from openkyrozen.providers.models import received_response, ModelResponse, model_response from openkyrozen.providers.config import ProviderConfig +from openkyrozen.providers.models import ( + ModelResponse, + model_response, + received_response, +) from openkyrozen.providers.retry import _retry_with_backoff class GoogleProvider(LLMProvider): - """Handles Google Gemini through the current Google Gen AI SDK.""" + """Handles Google Gemini through the Google Gen AI SDK.""" def __init__(self, config: ProviderConfig) -> None: super().__init__(config) + try: from google import genai except ImportError: - sys.exit( + raise ImportError( "The 'google-genai' package is required for Gemini. " "Install it with: pip install google-genai" + ) from None + + api_key = getattr(config, "api_key", None) or os.getenv( + "GEMINI_API_KEY" + ) or os.getenv("GOOGLE_API_KEY") + + if not api_key: + raise ValueError( + "Gemini API key is missing. Set GEMINI_API_KEY or " + "configure the provider API key." ) - self._client = genai.Client(api_key=config.api_key or os.environ.get("GEMINI_API_KEY", "")) + + self._client = genai.Client(api_key=api_key) @staticmethod - def _contents(messages: list[dict[str, str]]) -> tuple[list[dict[str, Any]], str | None]: + def _json_value(value: Any) -> Any: + """Convert values into JSON-compatible Python values.""" + if isinstance(value, Mapping): + return dict(value) + + if isinstance(value, (list, tuple)): + return list(value) + + if isinstance(value, str): + try: + return json.loads(value) + except (json.JSONDecodeError, TypeError): + return value + + if value is None: + return {} + + if isinstance(value, (bool, int, float)): + return value + + return str(value) + + @classmethod + def _function_args(cls, value: Any) -> dict[str, Any]: + """Gemini function-call args must be an object.""" + value = cls._json_value(value) + + if isinstance(value, dict): + return value + + if isinstance(value, str): + try: + parsed = json.loads(value) + except json.JSONDecodeError: + return {} + + if isinstance(parsed, dict): + return parsed + + return {} + + @classmethod + def _contents( + cls, + messages: list[dict[str, Any]], + ) -> tuple[list[dict[str, Any]], str | None]: + """ + Convert conversation messages into Gemini contents. + + Tool responses are only emitted as function_response parts when + a matching function call exists in the preceding assistant turn. + """ contents: list[dict[str, Any]] = [] - system: list[str] = [] + system_messages: list[str] = [] + + pending_calls: dict[str, str] = {} + for message in messages: - role = message.get("role", "user") - text = str(message.get("content", "")) + role = str(message.get("role", "user")).lower() + content = message.get("content", "") + name = message.get("name") or message.get("tool_name") + call_id = ( + message.get("tool_call_id") + or message.get("call_id") + or message.get("id") + ) + if role == "system": - system.append(text) + if content: + system_messages.append(str(content)) continue - contents.append({ - "role": "model" if role == "assistant" else "user", - "parts": [{"text": text}], - }) - return contents or [{"role": "user", "parts": [{"text": "Continue."}]}], ( - "\n\n".join(system) if system else None + + if role in ("tool", "function"): + payload = cls._json_value(content) + + # Match the receipt to the original assistant call. + function_name = pending_calls.pop(str(call_id), None) if call_id else None + + # Some internal message formats identify the tool by name. + if function_name is None and name: + function_name = str(name) + + if function_name is not None: + if not isinstance(payload, dict): + payload = {"result": payload} + + contents.append( + { + "role": "user", + "parts": [ + { + "function_response": { + "name": function_name, + "response": payload, + } + } + ], + } + ) + else: + # Never send an unmatched function_response. + # Preserve the receipt as ordinary text instead. + contents.append( + { + "role": "user", + "parts": [ + { + "text": ( + f"Tool result ({name or 'unknown tool'}): " + f"{json.dumps(payload, ensure_ascii=False)}" + ) + } + ], + } + ) + continue + + if role == "assistant": + parts: list[dict[str, Any]] = [] + + if isinstance(content, str) and content.strip(): + parts.append({"text": content}) + elif content is not None and not isinstance(content, str): + normalized = cls._json_value(content) + if normalized: + parts.append( + { + "text": json.dumps( + normalized, + ensure_ascii=False, + ) + } + ) + + calls = message.get("calls") or message.get("tool_calls") or [] + + for call in calls: + if isinstance(call, Mapping): + call_id = call.get("id") or call.get("call_id") + function = call.get("function", call) + function_name = function.get("name") + arguments = function.get( + "arguments", + function.get("args", {}), + ) + elif isinstance(call, (list, tuple)) and len(call) >= 3: + call_id, function_name, arguments = call[:3] + else: + continue + + if not function_name: + continue + + function_name = str(function_name) + + parts.append( + { + "function_call": { + "name": function_name, + "args": cls._function_args(arguments), + } + } + ) + + if call_id is not None: + pending_calls[str(call_id)] = function_name + + if parts: + contents.append({"role": "model", "parts": parts}) + continue + + # Default: user message. + if isinstance(content, str): + text = content + else: + text = json.dumps( + cls._json_value(content), + ensure_ascii=False, + ) + + contents.append( + { + "role": "user", + "parts": [{"text": text or " "}], + } + ) + + if not contents: + contents.append( + { + "role": "user", + "parts": [{"text": "Continue."}], + } + ) + + system_instruction = ( + "\n\n".join(system_messages) if system_messages else None ) + return contents, system_instruction + @staticmethod - def _usage(response: Any) -> dict[str, int | None] | None: - meta = getattr(response, "usage_metadata", None) - if meta is None: - return None - return { - "prompt_tokens": getattr(meta, "prompt_token_count", 0) or 0, - "completion_tokens": getattr(meta, "candidates_token_count", 0) or 0, - } - - def chat(self, messages: list[dict[str, str]], model: str | None = None) -> tuple[str, dict | None]: - return self.chat_response(messages, model).as_legacy_tuple() - - def chat_response(self, messages: list[dict[str, str]], model: str | None = None) -> ModelResponse: - model = model or self.config.model_simple - started = time.monotonic() + def _response_text(response: Any) -> str: + """Extract visible text from a Gemini response.""" + text = getattr(response, "text", None) + if isinstance(text, str): + return text + + output: list[str] = [] + + for candidate in getattr(response, "candidates", None) or []: + content = getattr(candidate, "content", None) + for part in getattr(content, "parts", None) or []: + part_text = getattr(part, "text", None) + if isinstance(part_text, str) and not getattr( + part, "thought", False + ): + output.append(part_text) + + return "".join(output) + + @staticmethod + def _response_calls(response: Any) -> list[tuple[str, str, dict[str, Any]]]: + """Extract function calls while preserving IDs when available.""" + calls: list[tuple[str, str, dict[str, Any]]] = [] + + for candidate in getattr(response, "candidates", None) or []: + content = getattr(candidate, "content", None) + + for part in getattr(content, "parts", None) or []: + function_call = getattr(part, "function_call", None) + if not function_call: + continue + + name = getattr(function_call, "name", None) + if not name: + continue + + args = getattr(function_call, "args", None) or {} + if not isinstance(args, dict): + args = GoogleProvider._function_args(args) + + call_id = getattr(part, "id", None) or "" + calls.append((str(call_id), str(name), args)) + + return calls + + def chat_response( + self, + messages: list[dict[str, Any]], + **kwargs: Any, + ) -> ModelResponse: + """Send a conversation to Gemini and return the provider response.""" + from google.genai import types contents, system_instruction = self._contents(messages) - request_config: dict[str, Any] = {} + + model = ( + kwargs.pop("model", None) + or getattr(self._config, "model", None) + or "gemini-2.5-flash" + ) + + config_kwargs: dict[str, Any] = {} if system_instruction: - request_config["system_instruction"] = system_instruction + config_kwargs["system_instruction"] = system_instruction - def _call(): - return self._client.models.generate_content( - model=model, contents=contents, config=request_config or None, - ) + generation_config = types.GenerateContentConfig(**config_kwargs) + + start = time.monotonic() + + response = self._client.models.generate_content( + model=model, + contents=contents, + config=generation_config, + ) + + elapsed = time.monotonic() - start + text = self._response_text(response) + calls = self._response_calls(response) + + usage = getattr(response, "usage_metadata", None) + if usage is not None: + usage_data = { + "prompt_tokens": getattr(usage, "prompt_token_count", 0) or 0, + "completion_tokens": getattr( + usage, "candidates_token_count", 0 + ) or 0, + "total_tokens": getattr(usage, "total_token_count", 0) or 0, + } + else: + usage_data = {} + + result = model_response( + text=text, + tool_calls=calls, + usage=usage_data, + elapsed=elapsed, + ) + + return result + + def stream_response( + self, + messages: list[dict[str, Any]], + **kwargs: Any, + ) -> Iterator[str]: + """Stream generated text from Gemini.""" + from google.genai import types - response = _retry_with_backoff(_call) - with received_response(): - usage_dict = self._usage(response) - usage_ledger._track_cost(self.config.provider, usage_dict, model=model, - latency_ms=round((time.monotonic() - started) * 1000)) - candidates = getattr(response, "candidates", None) or () - candidate = candidates[0] if candidates else None - parts = getattr(getattr(candidate, "content", None), "parts", None) - text = ("".join(part.text for part in parts if getattr(part, "text", None) - and not getattr(part, "thought", False)) if parts is not None - else str(getattr(response, "text", "") or "")) - calls = [] - for part in parts or (): - call = getattr(part, "function_call", None) - if call is not None: - arguments = getattr(call, "args", None) - calls.append((getattr(call, "id", None), getattr(call, "name", None), - {} if arguments is None else arguments)) - block_reason = getattr(getattr(response, "prompt_feedback", None), "block_reason", None) - return model_response(provider=self.name, model=model, actual_model=getattr(response, "model_version", None), text=text, usage=usage_dict, - calls=calls, response_id=getattr(response, "response_id", None), - raw_finish_reason=getattr(candidate, "finish_reason", None) or block_reason, - blocked=block_reason not in {None, "BLOCK_REASON_UNSPECIFIED", "BLOCKED_REASON_UNSPECIFIED"}) - - def chat_stream(self, messages: list[dict[str, str]], model: str | None = None) -> Iterator[str]: - model = model or self.config.model_simple contents, system_instruction = self._contents(messages) - request_config: dict[str, Any] = {} + + model = ( + kwargs.pop("model", None) + or getattr(self._config, "model", None) + or "gemini-2.5-flash" + ) + + config_kwargs: dict[str, Any] = {} if system_instruction: - request_config["system_instruction"] = system_instruction - started = time.monotonic() - final_usage: dict[str, int | None] | None = None - completed = False - stream = _retry_with_backoff(lambda: self._client.models.generate_content_stream( - model=model, contents=contents, config=request_config or None, - )) - try: - for chunk in stream: - usage = self._usage(chunk) - if usage is not None: - final_usage = usage - delta = str(getattr(chunk, "text", "") or "") - if delta: - yield delta - completed = True - finally: - if completed: - usage_ledger._track_cost(self.config.provider, final_usage, model=model, - latency_ms=round((time.monotonic() - started) * 1000)) - - -class VertexProvider(GoogleProvider): - """Google Gen AI SDK configured for Vertex AI and ADC.""" + config_kwargs["system_instruction"] = system_instruction - def __init__(self, config: ProviderConfig) -> None: - LLMProvider.__init__(self, config) - try: - from google import genai - except ImportError: - sys.exit( - "The 'google-genai' package is required for Vertex AI. " - "Install it with: pip install google-genai" - ) - self._client = genai.Client( - vertexai=True, - project=os.environ.get("GOOGLE_CLOUD_PROJECT"), - location=os.environ.get("GOOGLE_CLOUD_LOCATION", "global"), + config = types.GenerateContentConfig(**config_kwargs) + + stream = self._client.models.generate_content_stream( + model=model, + contents=contents, + config=config, ) + + for chunk in stream: + text = self._response_text(chunk) + if text: + yield text