diff --git a/README.md b/README.md index dfc70ad7..05f89b91 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ codewiki --version ### 2. Configure Your Environment -CodeWiki supports multiple LLM providers: **OpenAI-compatible**, **Anthropic**, **AWS Bedrock**, and **Azure OpenAI**. +CodeWiki supports multiple LLM providers: **OpenAI-compatible**, **Anthropic**, **AWS Bedrock**, **Azure OpenAI**, and **Google Gemini**. ```bash # Anthropic @@ -68,6 +68,13 @@ codewiki config set \ --aws-region us-east-1 \ --main-model anthropic.claude-sonnet-4-v2:0 \ --cluster-model anthropic.claude-sonnet-4-v2:0 + +# Google Gemini +codewiki config set \ + --provider gemini \ + --api-key YOUR_GEMINI_API_KEY \ + --main-model gemini-2.5-pro \ + --cluster-model gemini-2.5-pro ``` ### 3. Generate Documentation @@ -351,7 +358,7 @@ CodeWiki employs a three-stage process for comprehensive documentation generatio - **Python 3.12+** - **Node.js** (for Mermaid diagram validation) -- **LLM API access** (Anthropic Claude, OpenAI, Azure OpenAI, AWS Bedrock) +- **LLM API access** (Anthropic Claude, OpenAI, Azure OpenAI, AWS Bedrock, Google Gemini) - **Git** (for branch creation features) --- diff --git a/codewiki/cli/commands/config.py b/codewiki/cli/commands/config.py index 63df5609..49f85879 100644 --- a/codewiki/cli/commands/config.py +++ b/codewiki/cli/commands/config.py @@ -85,7 +85,7 @@ def config_group(): ) @click.option( "--provider", - type=click.Choice(['openai-compatible', 'anthropic', 'bedrock', 'azure-openai'], case_sensitive=False), + type=click.Choice(['openai-compatible', 'anthropic', 'bedrock', 'azure-openai', 'gemini'], case_sensitive=False), help="LLM provider type (default: openai-compatible)" ) @click.option( @@ -565,6 +565,11 @@ def config_validate(quick: bool, verbose: bool): azure_endpoint=config.base_url, ) client.models.list() + elif provider == "gemini": + # Use Google Gemini SDK + from google import genai + gemini_client = genai.Client(api_key=api_key) + list(gemini_client.models.list()) elif "api.anthropic.com" in base_url_lower: # Use Anthropic SDK for native Anthropic endpoints import anthropic diff --git a/codewiki/src/be/llm_services.py b/codewiki/src/be/llm_services.py index db3437a7..d2a89492 100644 --- a/codewiki/src/be/llm_services.py +++ b/codewiki/src/be/llm_services.py @@ -4,7 +4,7 @@ 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, gemini. """ import logging from openai.types import chat @@ -57,6 +57,7 @@ def _get_litellm_model_name(model_name: str, provider: str) -> str: For Bedrock, prefixes the model name with 'bedrock/' if not already prefixed. For Anthropic, prefixes with 'anthropic/' if not already prefixed. + For Gemini, prefixes with 'gemini/' if not already prefixed. """ if provider == "bedrock": if not model_name.startswith("bedrock/"): @@ -64,6 +65,9 @@ def _get_litellm_model_name(model_name: str, provider: str) -> str: elif provider == "anthropic": if not model_name.startswith("anthropic/"): return f"anthropic/{model_name}" + elif provider == "gemini": + if not model_name.startswith("gemini/"): + return f"gemini/{model_name}" return model_name @@ -154,8 +158,9 @@ def call_llm( """ Call LLM with the given prompt. - Supports openai-compatible, anthropic, and bedrock providers. + Supports openai-compatible, anthropic, bedrock, azure-openai, and gemini providers. For bedrock/anthropic, uses litellm to translate the API call. + For gemini, uses the google-genai SDK directly. Args: prompt: The prompt to send @@ -177,6 +182,9 @@ def call_llm( if provider == "azure-openai": return _call_llm_via_azure(prompt, config, model, temperature) + if provider == "gemini": + return _call_llm_via_gemini(prompt, config, model, temperature) + # Default: OpenAI-compatible client = create_openai_client(config) @@ -260,3 +268,39 @@ def _call_llm_via_azure( max_tokens=config.max_tokens, ) return response.choices[0].message.content + + +def _call_llm_via_gemini( + prompt: str, + config: Config, + model: str, + temperature: float = 0.0 +) -> str: + """ + Call LLM via Google Gemini using the google-genai SDK. + + Uses the google.genai client with the provided API key. + The model name may optionally include a 'gemini/' prefix, which is + stripped before passing to the SDK. + """ + from google import genai + from google.genai import types as genai_types + + api_key = config.llm_api_key + # Strip the 'gemini/' litellm prefix if present + sdk_model = model[len("gemini/"):] if model.startswith("gemini/") else model + + logger.debug("Calling Gemini model %s via google-genai SDK", sdk_model) + + client = genai.Client(api_key=api_key) + response = client.models.generate_content( + model=sdk_model, + contents=prompt, + config=genai_types.GenerateContentConfig( + temperature=temperature, + max_output_tokens=config.max_tokens, + ), + ) + if response.text is None: + raise ValueError(f"Gemini model {sdk_model} returned an empty response") + return response.text diff --git a/codewiki/src/config.py b/codewiki/src/config.py index 120ac2bd..9644c798 100644 --- a/codewiki/src/config.py +++ b/codewiki/src/config.py @@ -58,7 +58,7 @@ class Config: cluster_model: str fallback_model: str = FALLBACK_MODEL_1 # Provider configuration - provider: str = "openai-compatible" # openai-compatible, anthropic, bedrock, azure-openai + provider: str = "openai-compatible" # openai-compatible, anthropic, bedrock, azure-openai, gemini aws_region: str = "us-east-1" api_version: str = "2024-12-01-preview" # Azure OpenAI API version azure_deployment: str = "" # Azure OpenAI deployment name @@ -181,7 +181,7 @@ def from_cli( main_model: Primary model cluster_model: Clustering model fallback_model: Fallback model - provider: LLM provider type (openai-compatible, anthropic, bedrock, azure-openai) + provider: LLM provider type (openai-compatible, anthropic, bedrock, azure-openai, gemini) aws_region: AWS region for Bedrock provider api_version: Azure OpenAI API version azure_deployment: Azure OpenAI deployment name