added search models

src/config.py

@@ -35,8 +35,8 @@ DEFAULT_SAFETY_SETTINGS = [
     {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}
 ]
 
-# Supported Models (for /v1beta/models endpoint)
-SUPPORTED_MODELS = [
+# Base Models (without search variants)
+BASE_MODELS = [
     {
         "name": "models/gemini-2.5-pro-preview-05-06",
         "version": "001",
@@ -93,7 +93,7 @@ SUPPORTED_MODELS = [
         "name": "models/gemini-2.5-flash-preview-04-17",
         "version": "001",
         "displayName": "Gemini 2.5 Flash Preview 04-17",
-        "description": "Preview version of Gemini 2.5 Flash from May 20th",
+        "description": "Preview version of Gemini 2.5 Flash from April 17th",
         "inputTokenLimit": 1048576,
         "outputTokenLimit": 65535,
         "supportedGenerationMethods": ["generateContent", "streamGenerateContent"],
@@ -115,4 +115,33 @@ SUPPORTED_MODELS = [
         "topP": 0.95,
         "topK": 64
     }
-]
+]
+
+# Generate search variants for applicable models
+def _generate_search_variants():
+    """Generate search variants for models that support content generation."""
+    search_models = []
+    for model in BASE_MODELS:
+        # Only add search variants for models that support content generation
+        if "generateContent" in model["supportedGenerationMethods"]:
+            search_variant = model.copy()
+            search_variant["name"] = model["name"] + "-search"
+            search_variant["displayName"] = model["displayName"] + " with Google Search"
+            search_variant["description"] = model["description"] + " (includes Google Search grounding)"
+            search_models.append(search_variant)
+    return search_models
+
+# Supported Models (includes both base models and search variants)
+SUPPORTED_MODELS = BASE_MODELS + _generate_search_variants()
+
+# Helper function to get base model name from search variant
+def get_base_model_name(model_name):
+    """Convert search variant model name to base model name."""
+    if model_name.endswith("-search"):
+        return model_name[:-7]  # Remove "-search" suffix
+    return model_name
+
+# Helper function to check if model uses search grounding
+def is_search_model(model_name):
+    """Check if model name indicates search grounding should be enabled."""
+    return model_name.endswith("-search")

src/google_api_client.py

@@ -11,7 +11,7 @@ from google.auth.transport.requests import Request as GoogleAuthRequest
 
 from .auth import get_credentials, save_credentials, get_user_project_id, onboard_user
 from .utils import get_user_agent
-from .config import CODE_ASSIST_ENDPOINT, DEFAULT_SAFETY_SETTINGS
+from .config import CODE_ASSIST_ENDPOINT, DEFAULT_SAFETY_SETTINGS, get_base_model_name, is_search_model
 import asyncio
 
 
@@ -310,7 +310,15 @@ def build_gemini_payload_from_native(native_request: dict, model_from_path: str)
     native_request["generationConfig"]["thinkingConfig"]["includeThoughts"] = True
     native_request["generationConfig"]["thinkingConfig"]["thinkingBudget"] = -1
     
+    # Add Google Search grounding for search models
+    if is_search_model(model_from_path):
+        if "tools" not in native_request:
+            native_request["tools"] = []
+        # Add googleSearch tool if not already present
+        if not any(tool.get("googleSearch") for tool in native_request["tools"]):
+            native_request["tools"].append({"googleSearch": {}})
+    
     return {
-        "model": model_from_path,
+        "model": get_base_model_name(model_from_path),  # Use base model name for API call
         "request": native_request
     }

src/openai_transformers.py

@@ -8,7 +8,7 @@ import uuid
 from typing import Dict, Any
 
 from .models import OpenAIChatCompletionRequest, OpenAIChatCompletionResponse
-from .config import DEFAULT_SAFETY_SETTINGS
+from .config import DEFAULT_SAFETY_SETTINGS, is_search_model, get_base_model_name
 
 
 def openai_request_to_gemini(openai_request: OpenAIChatCompletionRequest) -> Dict[str, Any]:
@@ -91,12 +91,19 @@ def openai_request_to_gemini(openai_request: OpenAIChatCompletionRequest) -> Dic
         if openai_request.response_format.get("type") == "json_object":
             generation_config["responseMimeType"] = "application/json"
 
-    return {
+    # Build the request payload
+    request_payload = {
         "contents": contents,
         "generationConfig": generation_config,
         "safetySettings": DEFAULT_SAFETY_SETTINGS,
-        "model": openai_request.model
+        "model": get_base_model_name(openai_request.model)  # Use base model name for API call
     }
+    
+    # Add Google Search grounding for search models
+    if is_search_model(openai_request.model):
+        request_payload["tools"] = [{"googleSearch": {}}]
+    
+    return request_payload
 
 
 def gemini_response_to_openai(gemini_response: Dict[str, Any], model: str) -> Dict[str, Any]: