diff --git a/CHANGELOG.md b/CHANGELOG.md index 90c3bb5f..42ce22de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ uses [Semantic Versioning](https://semver.org/). ### Added +- **Google provider.** `google` (Gemini API key) and `google-vertex` (Vertex AI + with Google credentials: ADC or Workload Identity Federation, plus + `--google-project` and `--google-location`) run on Gemini without a base + URL. Agent loops use pydantic-ai's native `GoogleModel`, so Gemini 3 thought + signatures survive tool-call turns, and the cost summary prices Gemini runs + from genai-prices ([#20](https://github.com/FSoft-AI4Code/CodeWiki/issues/20)). - **Rust analyzer.** Structs, unions, enums, traits, free functions, and impl/trait methods become components, with trait-impl, field-type, struct-literal, and call edges. `#[cfg(test)]` modules are skipped, and diff --git a/README.md b/README.md index 95a6edfa..55cd282f 100644 --- a/README.md +++ b/README.md @@ -73,8 +73,8 @@ codewiki config set --provider claude-code \ --main-model claude-sonnet-4-6 --cluster-model claude-sonnet-4-6 ``` -Anthropic, Azure OpenAI, AWS Bedrock, Atlas Cloud, and Codex are also -supported. See [Providers and models](./guides/providers.md). +Anthropic, Azure OpenAI, AWS Bedrock, Google (Gemini API and Vertex AI), +Atlas Cloud, and Codex are also supported. See [Providers and models](./guides/providers.md). **3. Generate** diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py index 072bbd88..4e5fcfa3 100644 --- a/codewiki/cli/adapters/doc_generator.py +++ b/codewiki/cli/adapters/doc_generator.py @@ -151,6 +151,8 @@ def generate(self) -> DocumentationJob: fallback_model=self.config.get("fallback_model"), provider=self.config.get("provider", "openai-compatible"), aws_region=self.config.get("aws_region", "us-east-1"), + google_project=self.config.get("google_project", ""), + google_location=self.config.get("google_location", ""), max_tokens=self.config.get("max_tokens", 32768), max_token_per_module=self.config.get("max_token_per_module", 36369), max_token_per_leaf_module=self.config.get("max_token_per_leaf_module", 16000), diff --git a/codewiki/cli/commands/config.py b/codewiki/cli/commands/config.py index e8cd3d4e..e81ea3bd 100644 --- a/codewiki/cli/commands/config.py +++ b/codewiki/cli/commands/config.py @@ -69,12 +69,16 @@ def config_group(): "azure-openai", "claude-code", "codex", + "google", + "google-vertex", ], case_sensitive=False, ), help=( "LLM provider type (default: openai-compatible). " "Use 'atlas-cloud' for Atlas Cloud (base URL auto-set; reads ATLASCLOUD_API_KEY). " + "Use 'google' for the Gemini API (reads GEMINI_API_KEY / GOOGLE_API_KEY) or " + "'google-vertex' for Vertex AI with Google credentials (no API key). " "Use 'claude-code' or 'codex' to run on a CLI subscription instead of an API key." ), ) @@ -83,6 +87,17 @@ def config_group(): "--api-version", type=str, help="Azure OpenAI API version (default: 2024-12-01-preview)" ) @click.option("--azure-deployment", type=str, help="Azure OpenAI deployment name") +@click.option( + "--google-project", + type=str, + help="Google Cloud project for google-vertex (default: GOOGLE_CLOUD_PROJECT)", +) +@click.option( + "--google-location", + type=str, + help="Vertex AI location for google-vertex, e.g. global, eu, us " + "(default: GOOGLE_CLOUD_LOCATION, else global)", +) @click.option( "--use-gitignore/--no-gitignore", default=None, @@ -110,6 +125,8 @@ def config_set( aws_region: Optional[str] = None, api_version: Optional[str] = None, azure_deployment: Optional[str] = None, + google_project: Optional[str] = None, + google_location: Optional[str] = None, use_gitignore: Optional[bool] = None, prompt_caching: Optional[bool] = None, ): @@ -142,6 +159,16 @@ def config_set( # Subscription mode (Codex) $ codewiki config set --provider codex --main-model gpt-5.2-codex + \b + # Vertex AI (Google credentials, no API key or base URL) + $ codewiki config set --provider google-vertex --google-project my-proj \\ + --google-location eu --main-model gemini-3.8-flash --cluster-model gemini-3.8-flash + + \b + # Gemini API (key read from GEMINI_API_KEY if --api-key is omitted) + $ codewiki config set --provider google --main-model gemini-3.8-flash \\ + --cluster-model gemini-3.8-flash + \b # Update only API key $ codewiki config set --api-key sk-new-key @@ -181,6 +208,8 @@ def config_set( aws_region, api_version, azure_deployment, + google_project, + google_location, use_gitignore is not None, prompt_caching is not None, ] @@ -199,6 +228,20 @@ def config_set( if not api_key: api_key = os.getenv("ATLASCLOUD_API_KEY") + manager = ConfigManager() + manager.load() # Load existing config if present + + if provider and provider.lower() == "google" and not api_key: + api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") + previous = manager.get_config() + if not api_key and not (previous and previous.provider == "google"): + # The keyring has one key slot: reusing it would send another + # provider's key to Google. + raise ConfigurationError( + "--provider google needs a Gemini API key: pass --api-key or set " + "GEMINI_API_KEY (or GOOGLE_API_KEY)." + ) + # Validate inputs before saving validated_data = {} @@ -265,9 +308,11 @@ def config_set( if prompt_caching is not None: validated_data["prompt_caching"] = prompt_caching - # Create config manager and save - manager = ConfigManager() - manager.load() # Load existing config if present + if google_project is not None: + validated_data["google_project"] = google_project + + if google_location is not None: + validated_data["google_location"] = google_location manager.save( api_key=validated_data.get("api_key"), @@ -285,6 +330,8 @@ def config_set( aws_region=validated_data.get("aws_region"), api_version=validated_data.get("api_version"), azure_deployment=validated_data.get("azure_deployment"), + google_project=validated_data.get("google_project"), + google_location=validated_data.get("google_location"), use_gitignore=validated_data.get("use_gitignore"), prompt_caching=validated_data.get("prompt_caching"), ) @@ -352,6 +399,12 @@ def config_set( if azure_deployment: click.secho(f"✓ Azure Deployment: {azure_deployment}", fg="green") + if google_project: + click.secho(f"✓ Google project: {google_project}", fg="green") + + if google_location: + click.secho(f"✓ Google location: {google_location}", fg="green") + if use_gitignore is not None: click.secho(f"✓ Use gitignore: {use_gitignore}", fg="green") @@ -405,6 +458,9 @@ def config_show(output_json: bool): "api_key": mask_api_key(api_key) if api_key else "Not set", "api_key_storage": "keychain" if manager.keyring_available else "encrypted_file", "base_url": config.base_url if config else "", + "provider": config.provider if config else "", + "google_project": config.google_project if config else "", + "google_location": config.google_location if config else "", "main_model": config.main_model if config else "", "cluster_model": config.cluster_model if config else "", "fallback_model": config.fallback_model if config else "glm-4p5", @@ -430,12 +486,14 @@ def config_show(output_json: bool): click.echo("━" * 40) click.echo() - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import is_caw_provider, needs_base_url caw_mode = bool(config) and is_caw_provider(config.provider) click.secho("Credentials", fg="cyan", bold=True) - if caw_mode: + if config and config.provider == "google-vertex": + click.secho(" Google ADC / Workload Identity (no API key needed)", fg="cyan") + elif caw_mode: cli_name = "claude" if config.provider == "claude-code" else "codex" click.secho( f" Subscription mode: authenticate via '{cli_name} login' (no API key needed)", @@ -453,10 +511,18 @@ def config_show(output_json: bool): click.echo(f" Provider: {config.provider}") click.echo(f" Main Model: {config.main_model or 'Not set'}") if not caw_mode: - click.echo(f" Base URL: {config.base_url or 'Not set'}") + if needs_base_url(config.provider): + click.echo(f" Base URL: {config.base_url or 'Not set'}") click.echo(f" Cluster Model: {config.cluster_model or 'Not set'}") click.echo(f" Fallback Model: {config.fallback_model or 'Not set'}") - if config.provider == "bedrock": + if config.provider == "google-vertex": + project = config.google_project or "from GOOGLE_CLOUD_PROJECT" + location = config.google_location or ( + "from GOOGLE_CLOUD_LOCATION, else global" + ) + click.echo(f" Google Project: {project}") + click.echo(f" Google Location: {location}") + elif config.provider == "bedrock": click.echo(f" AWS Region: {config.aws_region}") elif config.provider == "azure-openai": click.echo(f" API Version: {config.api_version}") @@ -565,20 +631,29 @@ def config_validate(quick: bool, verbose: bool): # Load config early so we know the provider for the rest of the checks. config = manager.get_config() - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import ( + is_caw_provider, + is_google_provider, + needs_api_key, + needs_base_url, + ) caw_mode = bool(config) and is_caw_provider(config.provider) + key_needed = bool(config) and needs_api_key(config.provider) + url_needed = bool(config) and needs_base_url(config.provider) + api_key = None - # Step 2: Check API key (skipped for subscription providers) + # Step 2: Check API key (skipped for subscription providers and Vertex AI) if verbose: click.echo() click.echo("[2/5] Checking API key...") - if caw_mode: + if not key_needed: + reason = "subscription mode" if caw_mode else "Google credentials" if verbose: - click.secho(" ✓ API key not required (subscription mode)", fg="green") + click.secho(f" ✓ API key not required ({reason})", fg="green") else: - click.secho("✓ API key not required (subscription mode)", fg="green") + click.secho(f"✓ API key not required ({reason})", fg="green") else: if verbose: storage = "system keychain" if manager.keyring_available else "encrypted file" @@ -597,16 +672,17 @@ def config_validate(quick: bool, verbose: bool): else: click.secho("✓ API key present (stored in keychain)", fg="green") - # Step 3: Check base URL (skipped for subscription providers) + # Step 3: Check base URL (skipped for subscription and Google providers) if verbose: click.echo() click.echo("[3/5] Checking base URL...") - if caw_mode: + if not url_needed: + reason = "subscription mode" if caw_mode else "Google endpoint" if verbose: - click.secho(" ✓ Base URL not required (subscription mode)", fg="green") + click.secho(f" ✓ Base URL not required ({reason})", fg="green") else: - click.secho("✓ Base URL not required (subscription mode)", fg="green") + click.secho(f"✓ Base URL not required ({reason})", fg="green") else: if verbose: click.echo(f" URL: {config.base_url}") @@ -657,7 +733,7 @@ def config_validate(quick: bool, verbose: bool): # Warn about non-top-tier cluster model if not is_top_tier_model(config.cluster_model): click.secho( - "⚠️ Cluster model is not top-tier. Consider using claude-sonnet-4 or gpt-4.", + "⚠️ Cluster model is not top-tier. Consider a Claude Sonnet/Opus, GPT-5 or Gemini 3 model.", fg="yellow", ) @@ -690,6 +766,26 @@ def config_validate(quick: bool, verbose: bool): f"✓ {cli_name} CLI available (run '{cli_name} login' if not yet authenticated)", fg="green", ) + elif is_google_provider(config.provider) and not quick: + if verbose: + click.echo() + click.echo("[5/5] Testing Google API access...") + try: + from codewiki.src.be import google_provider + + google_provider.check_access( + config.provider, + api_key if key_needed else None, + config.main_model, + config.google_project, + config.google_location, + ) + click.secho("✓ API connectivity test successful", fg="green") + except Exception as e: + click.secho("✗ API connectivity test failed", fg="red") + if verbose: + click.echo(f" Error: {e}") + sys.exit(EXIT_CONFIG_ERROR) elif not quick: if verbose: click.echo() diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index d2122edc..5d8b74c9 100644 --- a/codewiki/cli/commands/generate.py +++ b/codewiki/cli/commands/generate.py @@ -793,6 +793,8 @@ def generate_command( "api_key": api_key, "provider": getattr(config, "provider", "openai-compatible"), "aws_region": getattr(config, "aws_region", "us-east-1"), + "google_project": getattr(config, "google_project", ""), + "google_location": getattr(config, "google_location", ""), "agent_instructions": agent_instructions_dict, # Max token settings (runtime overrides take precedence) "max_tokens": max_tokens if max_tokens is not None else config.max_tokens, diff --git a/codewiki/cli/config_manager.py b/codewiki/cli/config_manager.py index 15f9f820..1881a7b9 100644 --- a/codewiki/cli/config_manager.py +++ b/codewiki/cli/config_manager.py @@ -142,6 +142,8 @@ def save( aws_region: Optional[str] = None, api_version: Optional[str] = None, azure_deployment: Optional[str] = None, + google_project: Optional[str] = None, + google_location: Optional[str] = None, use_gitignore: Optional[bool] = None, prompt_caching: Optional[bool] = None, ): @@ -165,6 +167,8 @@ def save( aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name + google_project: Google Cloud project for google-vertex + google_location: Vertex AI location for google-vertex use_gitignore: Apply Git ignore rules during repository analysis prompt_caching: Add prompt-cache breakpoints to agentic LLM calls """ @@ -190,6 +194,10 @@ def save( agent_instructions=AgentInstructions(), ) + from codewiki.src.be.backend import is_caw_provider, is_google_provider, needs_base_url + + was_google = is_google_provider(self._config.provider) + # Update fields if provided if base_url is not None: self._config.base_url = base_url @@ -221,6 +229,10 @@ def save( self._config.api_version = api_version if azure_deployment is not None: self._config.azure_deployment = azure_deployment + if google_project is not None: + self._config.google_project = google_project + if google_location is not None: + self._config.google_location = google_location if use_gitignore is not None: self._config.use_gitignore = use_gitignore if prompt_caching is not None: @@ -230,12 +242,29 @@ def save( # Caw providers only need main_model; API providers need base_url + # cluster_model on top of that. The validate() method itself routes # by provider, so we only gate on whether enough is set to validate. - from codewiki.src.be.backend import is_caw_provider + + # The glm-4p5 default fallback does not exist on Google endpoints; a + # fallback that 404s only after the main model failed hides the real error. + if ( + is_google_provider(self._config.provider) + and self._config.main_model + and self._config.fallback_model in ("", "glm-4p5") + ): + self._config.fallback_model = self._config.main_model + # Leaving Google: the Gemini fallback set above would 404 on the new endpoint. + elif ( + was_google and not is_google_provider(self._config.provider) and fallback_model is None + ): + self._config.fallback_model = "glm-4p5" if is_caw_provider(self._config.provider): if self._config.main_model: self._config.validate() - elif self._config.base_url and self._config.main_model and self._config.cluster_model: + elif ( + (self._config.base_url or not needs_base_url(self._config.provider)) + and self._config.main_model + and self._config.cluster_model + ): self._config.validate() # Save API key to keyring, falling back to file @@ -296,7 +325,8 @@ def is_configured(self) -> bool: Check if configuration is complete and valid. Subscription-mode providers (claude-code, codex) do not require an - API key — they authenticate via the underlying CLI's OAuth. + API key — they authenticate via the underlying CLI's OAuth, and + google-vertex uses Google credentials. Returns: True if configured, False otherwise @@ -304,9 +334,9 @@ def is_configured(self) -> bool: if self._config is None: return False - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import needs_api_key - if not is_caw_provider(self._config.provider): + if needs_api_key(self._config.provider): # Check if API key is set if self.get_api_key() is None: return False diff --git a/codewiki/cli/models/config.py b/codewiki/cli/models/config.py index 0a08acfa..0d6be21a 100644 --- a/codewiki/cli/models/config.py +++ b/codewiki/cli/models/config.py @@ -151,6 +151,8 @@ class Configuration: aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name + google_project: Google Cloud project for google-vertex + google_location: Vertex AI location for google-vertex max_tokens: Maximum tokens for LLM response (default: 32768) max_token_per_module: Maximum tokens per module for clustering (default: 36369) max_token_per_leaf_module: Maximum tokens per leaf module (default: 16000) @@ -171,6 +173,8 @@ class Configuration: aws_region: str = "us-east-1" api_version: str = "2024-12-01-preview" azure_deployment: str = "" + google_project: str = "" + google_location: str = "" max_tokens: int = 32768 max_token_per_module: int = 36369 max_token_per_leaf_module: int = 16000 @@ -191,12 +195,13 @@ def validate(self): Raises: ConfigurationError: If validation fails """ - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import is_caw_provider, needs_base_url if is_caw_provider(self.provider): validate_model_name(self.main_model) return - validate_url(self.base_url) + if needs_base_url(self.provider): + validate_url(self.base_url) validate_model_name(self.main_model) validate_model_name(self.cluster_model) validate_model_name(self.fallback_model) @@ -212,6 +217,8 @@ def to_dict(self) -> dict: "aws_region": self.aws_region, "api_version": self.api_version, "azure_deployment": self.azure_deployment, + "google_project": self.google_project, + "google_location": self.google_location, "max_tokens": self.max_tokens, "max_token_per_module": self.max_token_per_module, "max_token_per_leaf_module": self.max_token_per_leaf_module, @@ -251,6 +258,8 @@ def from_dict(cls, data: dict) -> "Configuration": aws_region=data.get("aws_region", "us-east-1"), api_version=data.get("api_version", "2024-12-01-preview"), azure_deployment=data.get("azure_deployment", ""), + google_project=data.get("google_project", ""), + google_location=data.get("google_location", ""), max_tokens=data.get("max_tokens", 32768), max_token_per_module=data.get("max_token_per_module", 36369), max_token_per_leaf_module=data.get("max_token_per_leaf_module", 16000), @@ -267,15 +276,14 @@ def is_complete(self) -> bool: Subscription-mode providers (claude-code, codex) only require ``main_model``; ``base_url``, ``cluster_model`` and ``fallback_model`` - are unused. + are unused. Google providers need no ``base_url``. """ - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import is_caw_provider, needs_base_url if is_caw_provider(self.provider): return bool(self.main_model) - return bool( - self.base_url and self.main_model and self.cluster_model and self.fallback_model - ) + has_url = bool(self.base_url) or not needs_base_url(self.provider) + return bool(has_url and self.main_model and self.cluster_model and self.fallback_model) def to_backend_config( self, @@ -319,6 +327,8 @@ def to_backend_config( aws_region=self.aws_region, api_version=self.api_version, azure_deployment=self.azure_deployment, + google_project=self.google_project, + google_location=self.google_location, max_tokens=self.max_tokens, max_token_per_module=self.max_token_per_module, max_token_per_leaf_module=self.max_token_per_leaf_module, diff --git a/codewiki/cli/utils/validation.py b/codewiki/cli/utils/validation.py index 315aa9b2..ae1c86a2 100644 --- a/codewiki/cli/utils/validation.py +++ b/codewiki/cli/utils/validation.py @@ -241,6 +241,7 @@ def is_top_tier_model(model: str) -> bool: "gpt-4", "gpt-5", "gemini-2.5", + "gemini-3", ] model_lower = model.lower() diff --git a/codewiki/mcp/server.py b/codewiki/mcp/server.py index dc689479..e5c1c5ab 100644 --- a/codewiki/mcp/server.py +++ b/codewiki/mcp/server.py @@ -505,10 +505,9 @@ async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: config = manager.get_config() api_key = manager.get_api_key() - from codewiki.src.be.backend import is_caw_provider + from codewiki.src.be.backend import needs_api_key - caw_mode = bool(config) and is_caw_provider(getattr(config, "provider", "")) - if not api_key and not caw_mode: + if not api_key and needs_api_key(getattr(config, "provider", "")): return [ _text( json.dumps( @@ -547,6 +546,8 @@ async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: fallback_model=config.fallback_model, provider=getattr(config, "provider", "openai-compatible"), aws_region=getattr(config, "aws_region", "us-east-1"), + google_project=getattr(config, "google_project", ""), + google_location=getattr(config, "google_location", ""), max_tokens=config.max_tokens, agent_instructions=agent_instructions or None, use_gitignore=arguments.get("use_gitignore", True), diff --git a/codewiki/src/be/backend.py b/codewiki/src/be/backend.py index 8d1f7864..d2c20cd0 100644 --- a/codewiki/src/be/backend.py +++ b/codewiki/src/be/backend.py @@ -8,7 +8,9 @@ Two implementations satisfy this interface: * :class:`PydanticAIBackend` — wraps the existing openai-compatible / anthropic - / bedrock / azure-openai paths via pydantic-ai + litellm. API-key based. + / bedrock / azure-openai paths via pydantic-ai + litellm, and the native + Google paths (``google``, ``google-vertex``) via pydantic-ai + google-genai. + API-key based, except Vertex AI which uses Google credentials. * :class:`CawBackend` — routes through the ``claude`` or ``codex`` CLI via the :mod:`caw` library, using the user's OAuth subscription. No API key. @@ -157,6 +159,28 @@ def is_caw_provider(provider: str) -> bool: return provider in CAW_PROVIDERS +GOOGLE_PROVIDERS = frozenset({"google", "google-vertex"}) + + +def is_google_provider(provider: str) -> bool: + """Return True if *provider* talks to Gemini natively (Gemini API or Vertex AI).""" + return provider in GOOGLE_PROVIDERS + + +def needs_api_key(provider: str) -> bool: + """Return True if *provider* authenticates with a stored API key. + + caw providers use the CLI's OAuth login, google-vertex uses Google + credentials (ADC / Workload Identity Federation). + """ + return not is_caw_provider(provider) and provider != "google-vertex" + + +def needs_base_url(provider: str) -> bool: + """Return True if *provider* needs a configured base URL.""" + return not is_caw_provider(provider) and not is_google_provider(provider) + + class LLMBackend(abc.ABC): """Abstract LLM backend used by the documentation generator. diff --git a/codewiki/src/be/google_provider.py b/codewiki/src/be/google_provider.py new file mode 100644 index 00000000..1cbca525 --- /dev/null +++ b/codewiki/src/be/google_provider.py @@ -0,0 +1,212 @@ +"""Native Google provider: Gemini API (``google``) and Vertex AI (``google-vertex``). + +Agent loops use pydantic-ai's GoogleModel, which round-trips Gemini 3 thought +signatures across tool-call turns. The genai client is built here and passed +as ``GoogleProvider(client=...)``: the one constructor shared by pydantic-ai +1.x and 2.x (2.x moved Vertex AI to ``GoogleCloudProvider``), and it keeps +``GOOGLE_API_KEY`` / ``GOOGLE_GENAI_USE_VERTEXAI`` in the environment from +switching modes. + +One-shot completions use the synchronous genai client, because callers run +them from inside an active event loop. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Iterable, Optional + +from google import genai +from google.genai import types as genai_types +from pydantic_ai.messages import ModelResponse +from pydantic_ai.models.google import GoogleModel, GoogleModelSettings +from pydantic_ai.providers.google import GoogleProvider +from pydantic_ai.usage import RequestUsage + +logger = logging.getLogger(__name__) + +# gemini-3.x models are served from the global, us and eu locations; the SDK +# default (us-central1) does not serve them. +DEFAULT_VERTEX_LOCATION = "global" +# The OpenAI SDK's default request timeout; genai sets none of its own. +_TIMEOUT_MS = 600_000 +_RETRY_ATTEMPTS = 3 +_MODEL_PREFIXES = ("google/", "models/", "google-gla:", "google-vertex:", "google-cloud:") + + +def normalise_model_name(name: str) -> str: + """Strip routing prefixes (``google/``, ``google-gla:`` ...) the Google API rejects.""" + for prefix in _MODEL_PREFIXES: + if name.startswith(prefix): + return name[len(prefix) :] + return name + + +def _http_options() -> genai_types.HttpOptions: + options: dict[str, Any] = {"timeout": _TIMEOUT_MS} + retry_cls = getattr(genai_types, "HttpRetryOptions", None) + if retry_cls is not None: # absent in older google-genai releases + options["retry_options"] = retry_cls(attempts=_RETRY_ATTEMPTS) + return genai_types.HttpOptions(**options) + + +def create_client( + provider: str, api_key: str | None, project: str = "", location: str = "" +) -> genai.Client: + """Build a genai client for ``google`` (API key) or ``google-vertex`` (ADC).""" + if provider == "google-vertex": + project = project or os.getenv("GOOGLE_CLOUD_PROJECT", "") + if not project: + raise ValueError( + "google-vertex needs a Google Cloud project: pass --google-project " + "('codewiki config set --google-project ') or set GOOGLE_CLOUD_PROJECT." + ) + # Never an api_key: with one, genai switches to Vertex Express Mode, and + # the stored key may belong to another provider. + return genai.Client( + vertexai=True, + project=project, + location=location or os.getenv("GOOGLE_CLOUD_LOCATION") or DEFAULT_VERTEX_LOCATION, + http_options=_http_options(), + ) + if not api_key: + raise ValueError( + "google needs a Gemini API key: run 'codewiki config set --api-key ' " + "or set GEMINI_API_KEY before 'config set'." + ) + return genai.Client(vertexai=False, api_key=api_key, http_options=_http_options()) + + +def check_access( + provider: str, api_key: str | None, model: str, project: str = "", location: str = "" +) -> None: + """Raise if the credentials cannot see *model* (auth, project, location or name wrong).""" + client = create_client(provider, api_key, project, location) + client.models.get(model=normalise_model_name(model)) + + +def _client_for(config) -> genai.Client: + return create_client( + config.provider, + config.llm_api_key, + getattr(config, "google_project", ""), + getattr(config, "google_location", ""), + ) + + +# One-shot calls reuse a sync client per credential set. Agent-loop models get a +# fresh client (_client_for): their async transport is bound to the event loop. +_SYNC_CLIENTS: dict[tuple[str, str | None, str, str], genai.Client] = {} + + +def _sync_client_for(config) -> genai.Client: + key = ( + config.provider, + config.llm_api_key, + getattr(config, "google_project", ""), + getattr(config, "google_location", ""), + ) + client = _SYNC_CLIENTS.get(key) + if client is None: + client = _SYNC_CLIENTS[key] = _client_for(config) + return client + + +def create_google_model(config, model_name: str) -> GoogleModel: + """GoogleModel for the agent loop. No temperature and no thinking config: + the model defaults apply (Google advises temperature 1.0 for Gemini 3).""" + return GoogleModel( + normalise_model_name(model_name), + provider=GoogleProvider(client=_client_for(config)), + settings=GoogleModelSettings(max_tokens=config.max_tokens), + ) + + +def _price(usage: dict[str, Any], model: str) -> Optional[float]: + """USD price of one call from pydantic-ai's genai-prices data, or None if unknown.""" + response = ModelResponse( + parts=[], + usage=RequestUsage( + input_tokens=usage.get("input_tokens", 0), + output_tokens=usage.get("output_tokens", 0), + cache_read_tokens=usage.get("cache_read_tokens", 0), + ), + model_name=model, + provider_name="google", + ) + return _response_cost(response) + + +def _response_cost(response: ModelResponse) -> Optional[float]: + try: + return float(response.cost().total_price) + except Exception: # noqa: BLE001 - unknown model or old genai-prices; cost is optional + return None + + +def messages_cost(messages: Iterable[Any]) -> Optional[float]: + """Summed price of the ModelResponses of an agent run, or None if none could be priced.""" + prices = [ + p for m in messages if isinstance(m, ModelResponse) if (p := _response_cost(m)) is not None + ] + return sum(prices) if prices else None + + +def _usage_dict(metadata) -> dict[str, Any]: + if metadata is None: + return {"requests": 1} + input_tokens = (metadata.prompt_token_count or 0) + (metadata.tool_use_prompt_token_count or 0) + output_tokens = (metadata.candidates_token_count or 0) + (metadata.thoughts_token_count or 0) + usage: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, # thinking is billed as output + "requests": 1, + } + if metadata.cached_content_token_count: + usage["cache_read_tokens"] = metadata.cached_content_token_count + return usage + + +def call_google( + prompt: str, config, model: str, system_prompt: str | None = None +) -> tuple[Optional[str], Optional[dict[str, Any]]]: + """One-shot completion. Returns (text or None, usage dict). + + None mirrors ``_extract_content``: empty, truncated-before-text or blocked + output is logged and returned as None, never raised. + """ + name = normalise_model_name(model) + client = _sync_client_for(config) + response = client.models.generate_content( + model=name, + contents=prompt, + config=genai_types.GenerateContentConfig( + system_instruction=system_prompt or None, + max_output_tokens=config.max_tokens, + ), + ) + usage = _usage_dict(response.usage_metadata) + price = _price(usage, response.model_version or name) + if price is not None: + usage["cost_usd"] = price + + candidate = response.candidates[0] if response.candidates else None + if candidate is None: + feedback = response.prompt_feedback + reason = feedback.block_reason if feedback and feedback.block_reason else None + logger.warning("Gemini returned no candidates (model=%s, block_reason=%s).", name, reason) + return None, usage + finish = candidate.finish_reason + parts = candidate.content.parts if candidate.content and candidate.content.parts else [] + text = "".join(p.text for p in parts if p.text and not p.thought) or None + if finish == genai_types.FinishReason.MAX_TOKENS: + logger.warning( + "Gemini response from %s stopped at max_output_tokens (finish_reason=MAX_TOKENS); " + "output is truncated%s. Thinking tokens count against this limit.", + name, + " and empty" if text is None else "", + ) + elif text is None: + logger.warning("Gemini returned no text (model=%s, finish_reason=%s).", name, finish) + return text, usage diff --git a/codewiki/src/be/llm_services.py b/codewiki/src/be/llm_services.py index f3b4acb2..de189aaf 100644 --- a/codewiki/src/be/llm_services.py +++ b/codewiki/src/be/llm_services.py @@ -4,7 +4,8 @@ Includes a compatibility layer for OpenAI-compatible API proxies that may return slightly non-standard responses (e.g. choices[].index = None). -Supports multiple providers: openai-compatible, anthropic, bedrock, azure-openai. +Supports multiple providers: openai-compatible, anthropic, bedrock, azure-openai, +google (Gemini API) and google-vertex (Vertex AI). """ import inspect @@ -19,11 +20,13 @@ from pydantic_ai.exceptions import ModelAPIError except ImportError: # older releases (e.g. the pinned 1.0.6) only have ModelHTTPError ModelAPIError = ModelHTTPError +from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings from pydantic_ai.models.fallback import FallbackModel from pydantic_ai.providers.openai import OpenAIProvider from openai import OpenAI, BadRequestError +from codewiki.src.be.backend import is_google_provider from codewiki.src.config import Config logger = logging.getLogger(__name__) @@ -233,8 +236,12 @@ def _create_litellm_openai_client(config: Config) -> OpenAI: ) -def create_main_model(config: Config) -> CachingOpenAIModel: +def create_main_model(config: Config) -> Model: """Create the main LLM model from configuration.""" + if is_google_provider(getattr(config, "provider", "")): + from codewiki.src.be.google_provider import create_google_model + + return create_google_model(config, config.main_model) return CachingOpenAIModel( model_name=config.main_model, prompt_caching=config.prompt_caching, @@ -244,8 +251,12 @@ def create_main_model(config: Config) -> CachingOpenAIModel: ) -def create_fallback_model(config: Config) -> CachingOpenAIModel: +def create_fallback_model(config: Config) -> Model: """Create the fallback LLM model from configuration.""" + if is_google_provider(getattr(config, "provider", "")): + from codewiki.src.be.google_provider import create_google_model + + return create_google_model(config, config.fallback_model) return CachingOpenAIModel( model_name=config.fallback_model, prompt_caching=config.prompt_caching, @@ -329,8 +340,9 @@ def call_llm( """ Call LLM with the given prompt. - Supports openai-compatible, anthropic, and bedrock providers. - For bedrock/anthropic, uses litellm to translate the API call. + Supports openai-compatible, anthropic, bedrock, azure-openai and google / + google-vertex providers. For bedrock/anthropic, uses litellm to translate + the API call; Google providers use the google-genai client. No temperature is sent; requests rely on the provider default (some models reject explicit values). @@ -350,6 +362,12 @@ def call_llm( provider = getattr(config, "provider", "openai-compatible") + if is_google_provider(provider): + from codewiki.src.be.google_provider import call_google + + text, _LAST_USAGE["usage"] = call_google(prompt, config, model, system_prompt) + return text + if provider in ("bedrock", "anthropic"): return _call_llm_via_litellm(prompt, config, model, system_prompt) diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py index 0aabd7f4..a945b382 100644 --- a/codewiki/src/be/pydantic_ai_backend.py +++ b/codewiki/src/be/pydantic_ai_backend.py @@ -23,7 +23,7 @@ ) from codewiki.src.be.agent_tools.read_code_components import read_code_components_tool from codewiki.src.be.agent_tools.str_replace_editor import str_replace_editor_tool -from codewiki.src.be.backend import AgentReply, LLMBackend, usage_to_dict +from codewiki.src.be.backend import AgentReply, LLMBackend, is_google_provider, usage_to_dict from codewiki.src.be.dependency_analyzer.models.core import Node from codewiki.src.be.llm_services import call_llm, create_fallback_models, pop_last_usage from codewiki.src.be.prompt_template import ( @@ -51,8 +51,24 @@ def _run_usage(result: Any) -> dict[str, Any] | None: return None +def _run_usage_with_cost(result: Any, provider: str) -> dict[str, Any] | None: + """`_run_usage`, plus ``cost_usd`` for Google providers (priced from genai-prices).""" + usage = _run_usage(result) + if usage is None or not is_google_provider(provider): + return usage + from codewiki.src.be.google_provider import messages_cost + + try: + cost = messages_cost(result.all_messages()) + except Exception: # noqa: BLE001 — cost is optional telemetry + cost = None + if cost is not None: + usage["cost_usd"] = cost + return usage + + class PydanticAIBackend(LLMBackend): - """API-key based backend using pydantic-ai + openai/litellm clients.""" + """API-key based backend using pydantic-ai + openai/litellm/google-genai clients.""" def __init__(self, config: Config) -> None: self._config = config @@ -103,7 +119,7 @@ async def run_update_agent( usage = self._record_usage( "update_agent", deps.current_module_name, - _run_usage(result), + _run_usage_with_cost(result, self._config.provider), seconds=seconds, model=self._config.main_model, ) @@ -198,7 +214,7 @@ async def run_module_agent( self._record_usage( "module_agent", module_name, - _run_usage(result), + _run_usage_with_cost(result, config.provider), seconds=time.time() - started, model=config.main_model, ) diff --git a/codewiki/src/config.py b/codewiki/src/config.py index 12176a45..0f0d88a6 100644 --- a/codewiki/src/config.py +++ b/codewiki/src/config.py @@ -93,11 +93,16 @@ class Config: fallback_model: str = FALLBACK_MODEL_1 # Provider configuration provider: str = ( - "openai-compatible" # openai-compatible, atlas-cloud, anthropic, bedrock, azure-openai + "openai-compatible" # openai-compatible, atlas-cloud, anthropic, bedrock, azure-openai, + # google, google-vertex ) aws_region: str = "us-east-1" api_version: str = "2024-12-01-preview" # Azure OpenAI API version azure_deployment: str = "" # Azure OpenAI deployment name + # Vertex AI (google-vertex): project falls back to GOOGLE_CLOUD_PROJECT, location to + # GOOGLE_CLOUD_LOCATION, then "global" + google_project: str = "" + google_location: str = "" # Max token settings max_tokens: int = DEFAULT_MAX_TOKENS max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE @@ -241,6 +246,8 @@ def from_cli( aws_region: str = "us-east-1", api_version: str = "2024-12-01-preview", azure_deployment: str = "", + google_project: str = "", + google_location: str = "", max_tokens: int = DEFAULT_MAX_TOKENS, max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE, max_token_per_leaf_module: int = DEFAULT_MAX_TOKEN_PER_LEAF_MODULE, @@ -268,10 +275,13 @@ def from_cli( main_model: Primary model cluster_model: Clustering model fallback_model: Fallback model - provider: LLM provider type (openai-compatible, atlas-cloud, anthropic, bedrock, azure-openai) + provider: LLM provider type (openai-compatible, atlas-cloud, anthropic, bedrock, + azure-openai, google, google-vertex) aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name + google_project: Google Cloud project for the google-vertex provider + google_location: Vertex AI location for the google-vertex provider max_tokens: Maximum tokens for LLM response max_token_per_module: Maximum tokens per module for clustering max_token_per_leaf_module: Maximum tokens per leaf module @@ -312,6 +322,8 @@ def from_cli( aws_region=aws_region, api_version=api_version, azure_deployment=azure_deployment, + google_project=google_project, + google_location=google_location, max_tokens=max_tokens, max_token_per_module=max_token_per_module, max_token_per_leaf_module=max_token_per_leaf_module, diff --git a/guides/cli-reference.md b/guides/cli-reference.md index d5b76c19..6f2043ae 100644 --- a/guides/cli-reference.md +++ b/guides/cli-reference.md @@ -166,7 +166,7 @@ keys you pass are changed. Provider examples are in [Providers](providers.md). | Flag | Meaning | | --- | --- | -| `--provider NAME` | One of `openai-compatible` (default), `atlas-cloud`, `anthropic`, `bedrock`, `azure-openai`, `claude-code`, `codex` | +| `--provider NAME` | One of `openai-compatible` (default), `atlas-cloud`, `anthropic`, `bedrock`, `azure-openai`, `claude-code`, `codex`, `google`, `google-vertex` | | `--api-key KEY` | API key. Stored in the system keychain when one is available | | `--base-url URL` | Provider endpoint. Set automatically for `atlas-cloud` | | `--main-model NAME` | Model for module documentation | @@ -175,6 +175,8 @@ keys you pass are changed. Provider examples are in [Providers](providers.md). | `--aws-region REGION` | Bedrock only | | `--api-version VERSION` | Azure OpenAI only | | `--azure-deployment NAME` | Azure OpenAI only | +| `--google-project ID` | `google-vertex` only. Defaults to `GOOGLE_CLOUD_PROJECT` | +| `--google-location LOC` | `google-vertex` only, e.g. `global`, `eu`, `us`. Defaults to `GOOGLE_CLOUD_LOCATION`, then `global` | | `--max-tokens N`, `--max-token-per-module N`, `--max-token-per-leaf-module N`, `--max-depth N` | Stored token limits, see the table above | | `--use-gitignore` / `--no-gitignore` | Stored default for Git ignore handling | | `--prompt-caching` / `--no-prompt-caching` | Stored default for prompt caching | diff --git a/guides/development.md b/guides/development.md index e20e8d8e..c6bf0b82 100644 --- a/guides/development.md +++ b/guides/development.md @@ -51,6 +51,7 @@ codewiki/ │ │ ├── agent_tools/ # read_code_components, str_replace_editor, sub-module delegation │ │ ├── backend.py # LLMBackend abstraction + factory │ │ ├── pydantic_ai_backend.py # API-key providers +│ │ ├── google_provider.py # google / google-vertex (Gemini API, Vertex AI) │ │ ├── caw_backend.py, caw_toolkit.py # claude-code / codex subscription providers │ │ ├── prompt_template.py, llm_services.py, module_naming.py │ └── fe/ # FastAPI web app (used by the Docker image) diff --git a/guides/providers.md b/guides/providers.md index aa6d3913..56f0c333 100644 --- a/guides/providers.md +++ b/guides/providers.md @@ -1,8 +1,8 @@ # Providers and models CodeWiki calls an LLM for module clustering and for writing each page. You -pick where those calls go with `codewiki config set --provider ...`. Seven -providers are supported. Two of them need no API key at all. +pick where those calls go with `codewiki config set --provider ...`. Nine +providers are supported. Three of them need no API key at all. | Provider | Needs | Good for | | --- | --- | --- | @@ -11,6 +11,8 @@ providers are supported. Two of them need no API key at all. | `anthropic` | API key | Direct Anthropic API | | `azure-openai` | API key + resource URL + deployment | Azure-hosted OpenAI models | | `bedrock` | AWS credentials + region | Anthropic and other models on AWS Bedrock | +| `google` | Gemini API key | Gemini models through the Gemini API (AI Studio) | +| `google-vertex` | Google credentials + project | Gemini models on Vertex AI, billed to a Google Cloud project, no API key | | `claude-code` | Claude Code CLI logged in | Run on a Claude Pro or Max subscription, no per-token billing | | `codex` | Codex CLI logged in | Run on a Codex subscription, no per-token billing | @@ -54,6 +56,46 @@ codewiki config set \ --cluster-model anthropic.claude-sonnet-4-v2:0 ``` +### Google (Gemini API and Vertex AI) + +Both Google providers talk to Gemini natively, so Gemini 3 thought signatures +survive tool-call turns. No base URL is needed. + +```bash +# Vertex AI: Google credentials (gcloud auth application-default login, or +# Workload Identity Federation in CI), no API key +codewiki config set \ + --provider google-vertex \ + --google-project YOUR_PROJECT \ + --google-location eu \ + --main-model gemini-3.8-flash \ + --cluster-model gemini-3.8-flash + +# Gemini API: the key is read from GEMINI_API_KEY (or GOOGLE_API_KEY) when +# --api-key is omitted +codewiki config set \ + --provider google \ + --main-model gemini-3.8-flash \ + --cluster-model gemini-3.8-flash +``` + +Notes: + +- **Location.** `--google-location` accepts `global`, `us` or `eu`; it defaults + to `GOOGLE_CLOUD_LOCATION`, then `global`. `eu` and `us` need google-genai + 2.x. `us-central1` does not serve Gemini 3.x. +- **Project.** `--google-project` defaults to `GOOGLE_CLOUD_PROJECT`. With + neither set, a run fails at startup. +- **Fallback model.** It defaults to the main model, because the usual default + (`glm-4p5`) does not exist on Google endpoints. +- **Output budget.** Thinking tokens share the `--max-tokens` budget. If pages + come back empty with `finish_reason=MAX_TOKENS`, raise `--max-tokens` to + 65536. +- **Privacy.** Google may use prompts sent with a free-tier AI Studio key for + product improvement. Use a paid key, or Vertex AI, for private code. +- **Cost.** The cost summary prices Gemini runs from the installed + `genai-prices` data. A model it does not know is reported as unpriced. + ### Atlas Cloud [Atlas Cloud](https://www.atlascloud.ai) is an inference platform that exposes diff --git a/pyproject.toml b/pyproject.toml index e101fbb2..c1a76ca0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dependencies = [ "pydantic>=2.11.7", "pydantic-settings>=2.10.1", "pydantic-ai>=1.0.6", + "google-genai>=1.36.0", "requests>=2.32.4", "python-dotenv>=1.1.1", "pathspec>=0.12.1", diff --git a/tests/test_google_config.py b/tests/test_google_config.py new file mode 100644 index 00000000..66ff6386 --- /dev/null +++ b/tests/test_google_config.py @@ -0,0 +1,251 @@ +"""Google (Gemini API / Vertex AI) provider: config, gates and CLI diagnostics.""" + +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner + +import codewiki.cli.commands.generate as gen +import codewiki.cli.config_manager as cm +from codewiki.cli.commands.config import config_set, config_show, config_validate +from codewiki.cli.models.config import Configuration +from codewiki.cli.utils.validation import is_top_tier_model +from codewiki.src.config import Config + +MODEL = "gemini-3.8-flash" + + +@pytest.fixture +def tmp_config(tmp_path, monkeypatch): + monkeypatch.setattr(cm, "CONFIG_DIR", tmp_path) + monkeypatch.setattr(cm, "CONFIG_FILE", tmp_path / "config.json") + monkeypatch.setattr(cm, "CREDENTIALS_FILE", tmp_path / "credentials.json") + monkeypatch.setenv("CODEWIKI_NO_KEYRING", "1") + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + return tmp_path + + +def _loaded_manager(): + manager = cm.ConfigManager() + manager.load() + return manager + + +def test_vertex_configuration_valid_without_url(): + cfg = Configuration( + provider="google-vertex", + base_url="", + main_model=MODEL, + cluster_model=MODEL, + fallback_model=MODEL, + google_project="test-proj", + google_location="eu", + ) + cfg.validate() + assert cfg.is_complete() is True + + again = Configuration.from_dict(cfg.to_dict()) + assert (again.google_location, again.google_project) == ("eu", "test-proj") + + +def test_vertex_is_configured_without_api_key_and_fallback_follows_main(tmp_config): + cm.ConfigManager().save( + provider="google-vertex", + main_model=MODEL, + cluster_model=MODEL, + google_project="test-proj", + ) + manager = _loaded_manager() + assert manager.is_configured() is True + assert manager.get_api_key() is None + assert manager.get_config().fallback_model == MODEL + + +def test_config_set_google_reads_gemini_api_key_env(tmp_config, monkeypatch): + monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-key") + result = CliRunner().invoke( + config_set, + ["--provider", "google", "--main-model", MODEL, "--cluster-model", MODEL], + ) + assert result.exit_code == 0, result.output + assert _loaded_manager().get_api_key() == "test-gemini-key" + + +def test_config_set_google_refuses_foreign_key(tmp_config): + cm.ConfigManager().save( + api_key="sk-old", + base_url="https://api.example.com/v1", + main_model="m", + cluster_model="m", + provider="openai-compatible", + ) + result = CliRunner().invoke(config_set, ["--provider", "google", "--main-model", MODEL]) + assert result.exit_code != 0 + assert "GEMINI_API_KEY" in result.output + assert _loaded_manager().get_config().provider == "openai-compatible" + + +def test_google_settings_reach_backend_config(tmp_path, monkeypatch): + seen = {} + cfg = SimpleNamespace( + main_model=MODEL, + cluster_model=MODEL, + fallback_model=MODEL, + base_url="", + provider="google-vertex", + google_project="test-proj", + google_location="eu", + aws_region="us-east-1", + agent_instructions=None, + max_tokens=1, + max_token_per_module=1, + max_token_per_leaf_module=1, + max_depth=1, + request_limit=1, + agent_retries=1, + use_gitignore=True, + prompt_caching=False, + ) + + class Stop(Exception): + pass + + class FakeManager: + def load(self): + return True + + def is_configured(self): + return True + + def get_config(self): + return cfg + + def get_api_key(self): + return None + + def fake_generator(**kwargs): + seen.update(kwargs["config"]) + raise Stop + + monkeypatch.setattr(gen, "ConfigManager", FakeManager) + monkeypatch.setattr(gen, "CLIDocumentationGenerator", fake_generator) + (tmp_path / "a.py").write_text("def f():\n return 1\n") + monkeypatch.chdir(tmp_path) + + CliRunner().invoke(gen.generate_command, ["-o", "out"]) + assert seen["google_project"] == "test-proj" + assert seen["google_location"] == "eu" + + backend = Config.from_cli( + repo_path=".", + output_dir="out", + llm_base_url="", + llm_api_key="", + main_model=MODEL, + cluster_model=MODEL, + provider="google-vertex", + google_project="test-proj", + google_location="eu", + ) + assert (backend.google_project, backend.google_location) == ("test-proj", "eu") + + +@pytest.fixture +def access_checks(monkeypatch): + """Record (provider, model) of every connectivity probe instead of calling Google.""" + from codewiki.src.be import google_provider + + calls = [] + + def fake(provider, api_key, model, project="", location=""): + calls.append((provider, model)) + + monkeypatch.setattr(google_provider, "check_access", fake) + return calls + + +def _save_vertex(): + cm.ConfigManager().save( + provider="google-vertex", + main_model=MODEL, + cluster_model=MODEL, + google_project="test-proj", + google_location="eu", + ) + + +def test_validate_vertex_without_key_or_url(tmp_config, access_checks): + _save_vertex() + result = CliRunner().invoke(config_validate, []) + assert result.exit_code == 0, result.output + assert "API key not required" in result.output + assert "Configuration is valid" in result.output + assert access_checks == [("google-vertex", MODEL)] + + +def test_validate_google_does_not_use_openai_client(tmp_config, access_checks, monkeypatch): + import openai + + def forbidden(*args, **kwargs): + raise AssertionError("a Gemini key must not reach an OpenAI client") + + monkeypatch.setattr(openai, "OpenAI", forbidden) + cm.ConfigManager().save( + api_key="test-gemini-key", + base_url="https://api.openai.com/v1", + provider="google", + main_model=MODEL, + cluster_model=MODEL, + ) + result = CliRunner().invoke(config_validate, []) + assert result.exit_code == 0, result.output + assert access_checks == [("google", MODEL)] + + +def test_validate_reports_google_access_failure(tmp_config, monkeypatch): + from codewiki.src.be import google_provider + + def failing(*args, **kwargs): + raise RuntimeError("404 model not found") + + monkeypatch.setattr(google_provider, "check_access", failing) + _save_vertex() + result = CliRunner().invoke(config_validate, []) + assert result.exit_code != 0 + assert "API connectivity test failed" in result.output + + +def test_show_vertex_lists_project_and_location(tmp_config): + import json + + _save_vertex() + result = CliRunner().invoke(config_show, []) + assert result.exit_code == 0, result.output + assert "Google ADC" in result.output + assert "test-proj" in result.output + assert "eu" in result.output + assert "Base URL" not in result.output + + as_json = CliRunner().invoke(config_show, ["--json"]) + assert json.loads(as_json.output)["google_project"] == "test-proj" + + +def test_gemini_3_is_top_tier(): + assert is_top_tier_model(MODEL) is True + + +def test_leaving_google_resets_gemini_fallback(tmp_config): + manager = cm.ConfigManager() + manager.save(provider="google-vertex", main_model=MODEL, cluster_model=MODEL) + assert manager.get_config().fallback_model == MODEL + + manager.save(provider="openai-compatible", base_url="https://api.example.com/v1") + assert manager.get_config().fallback_model == "glm-4p5" + + +def test_leaving_google_keeps_explicit_fallback(tmp_config): + manager = cm.ConfigManager() + manager.save(provider="google-vertex", main_model=MODEL, cluster_model=MODEL) + manager.save(provider="openai-compatible", fallback_model="gpt-5-mini") + assert manager.get_config().fallback_model == "gpt-5-mini" diff --git a/tests/test_google_provider.py b/tests/test_google_provider.py new file mode 100644 index 00000000..8a56812a --- /dev/null +++ b/tests/test_google_provider.py @@ -0,0 +1,229 @@ +"""Native Google provider: model factories, one-shot calls, usage and cost.""" + +import logging +from decimal import Decimal +from types import SimpleNamespace + +import pytest +from google.genai import types as gt +from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart +from pydantic_ai.models.fallback import FallbackModel +from pydantic_ai.models.google import GoogleModel +from pydantic_ai.usage import RequestUsage + +from codewiki.src.be import google_provider +from codewiki.src.be.llm_services import ( + call_llm, + create_fallback_models, + create_main_model, + pop_last_usage, +) +from codewiki.src.config import Config + +MODEL = "gemini-3.8-flash" + + +def _config(**overrides) -> Config: + kwargs = dict( + repo_path=".", + output_dir="out", + llm_base_url="", + llm_api_key="sk-1234", + main_model=MODEL, + cluster_model=MODEL, + fallback_model=MODEL, + provider="google-vertex", + google_project="test-proj", + google_location="eu", + ) + kwargs.update(overrides) + return Config.from_cli(**kwargs) + + +@pytest.fixture +def clean_env(monkeypatch): + for name in ( + "GOOGLE_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_GENAI_USE_VERTEXAI", + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_LOCATION", + ): + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def client_spy(monkeypatch): + calls = [] + real = google_provider.genai.Client + + def spy(**kwargs): + calls.append(kwargs) + return real(**kwargs) + + monkeypatch.setattr(google_provider.genai, "Client", spy) + return calls + + +def test_vertex_client_gets_no_api_key(clean_env, client_spy, monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "test-env-key") + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "false") + + model = create_main_model(_config()) + + assert len(client_spy) == 1 + kwargs = client_spy[0] + assert kwargs["vertexai"] is True + assert kwargs["project"] == "test-proj" + assert kwargs["location"] == "eu" + assert "api_key" not in kwargs + assert isinstance(model, GoogleModel) + assert model.client.vertexai is True + + +def test_vertex_location_defaults_to_global_and_project_is_required(clean_env, client_spy): + create_main_model(_config(google_location="")) + assert client_spy[-1]["location"] == "global" + + with pytest.raises(ValueError) as err: + create_main_model(_config(google_project="")) + assert "--google-project" in str(err.value) + assert "GOOGLE_CLOUD_PROJECT" in str(err.value) + + +def test_google_fallback_models_are_google_models(clean_env): + config = _config( + provider="google", + llm_api_key="test-key", + main_model="google/gemini-3.8-flash", + fallback_model="gemini-3.8-flash", + ) + chain = create_fallback_models(config) + + assert isinstance(chain, FallbackModel) + assert all(isinstance(m, GoogleModel) for m in chain.models) + assert chain.models[0].model_name == "gemini-3.8-flash" + + +def _response(candidates, **kwargs) -> gt.GenerateContentResponse: + return gt.GenerateContentResponse( + candidates=candidates, + usage_metadata=gt.GenerateContentResponseUsageMetadata( + prompt_token_count=100, + cached_content_token_count=40, + candidates_token_count=20, + thoughts_token_count=30, + ), + model_version=MODEL, + **kwargs, + ) + + +def _text_candidate(text: str, reason=gt.FinishReason.STOP) -> gt.Candidate: + return gt.Candidate( + content=gt.Content(role="model", parts=[gt.Part(text=text)]), finish_reason=reason + ) + + +class FakeModels: + def __init__(self, response): + self.response = response + self.calls = [] + + def generate_content(self, **kwargs): + self.calls.append(kwargs) + return self.response + + +@pytest.fixture +def fake_models(monkeypatch): + def install(response): + models = FakeModels(response) + monkeypatch.setattr(google_provider, "_SYNC_CLIENTS", {}) + monkeypatch.setattr( + google_provider, "_client_for", lambda config: SimpleNamespace(models=models) + ) + return models + + return install + + +def test_call_llm_google_text_and_usage(fake_models, monkeypatch): + models = fake_models(_response([_text_candidate("hello")])) + monkeypatch.setattr(google_provider, "_price", lambda usage, model: 0.25) + pop_last_usage() + + text = call_llm("p", _config(provider="google", llm_api_key="test-key"), system_prompt="s") + + assert text == "hello" + call = models.calls[0] + assert call["model"] == MODEL + assert call["config"].system_instruction == "s" + assert pop_last_usage() == { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_tokens": 40, + "requests": 1, + "cost_usd": 0.25, + } + + +@pytest.mark.parametrize( + "response, reason", + [ + ( + _response([gt.Candidate(content=None, finish_reason=gt.FinishReason.MAX_TOKENS)]), + "MAX_TOKENS", + ), + ( + gt.GenerateContentResponse( + candidates=None, + prompt_feedback=gt.GenerateContentResponsePromptFeedback( + block_reason=gt.BlockedReason.SAFETY + ), + ), + "SAFETY", + ), + ], + ids=["max-tokens", "blocked"], +) +def test_call_llm_google_empty_returns_none(fake_models, caplog, response, reason): + fake_models(response) + pop_last_usage() + + with caplog.at_level(logging.WARNING): + text = call_llm("p", _config(provider="google", llm_api_key="test-key")) + + assert text is None + assert any(r.levelno == logging.WARNING and reason in r.getMessage() for r in caplog.records) + if reason == "MAX_TOKENS": + assert pop_last_usage()["output_tokens"] == 50 + + +def test_cost_unknown_model_is_omitted(fake_models, monkeypatch): + response = _response([_text_candidate("hello")]) + response.model_version = "not-a-real-model" + fake_models(response) + pop_last_usage() + + assert call_llm("p", _config(provider="google", llm_api_key="test-key")) == "hello" + assert "cost_usd" not in pop_last_usage() + + priced = ModelResponse(parts=[], usage=RequestUsage(), model_name=MODEL) + monkeypatch.setattr( + ModelResponse, "cost", lambda self: SimpleNamespace(total_price=Decimal("0.5")) + ) + messages = [priced, ModelRequest(parts=[UserPromptPart("x")]), priced] + assert google_provider.messages_cost(messages) == 1.0 + + +def test_call_google_reuses_sync_client(monkeypatch): + built = [] + monkeypatch.setattr(google_provider, "_SYNC_CLIENTS", {}) + monkeypatch.setattr(google_provider, "_client_for", lambda config: built.append(1) or object()) + config = SimpleNamespace( + provider="google-vertex", llm_api_key=None, google_project="p", google_location="eu" + ) + first = google_provider._sync_client_for(config) + assert google_provider._sync_client_for(config) is first + assert len(built) == 1