feat: add genai-migrate codemod CLI for legacy google-generativeai

README.md

@@ -44,6 +44,51 @@ pip install google-genai
 uv pip install google-genai
 ```
 
+## Migrate from google-generativeai
+
+A codemod tool automates migration from the legacy `google-generativeai` SDK to `google-genai`. Install with the migrate extra and run on your codebase:
+
+```sh
+pip install google-genai[migrate]
+```
+
+```sh
+# Preview changes (unified diff)
+genai-migrate ./my_app/ --diff
+
+# Apply changes in-place
+genai-migrate ./my_app/ --in-place
+```
+
+**Before (legacy):**
+```python
+import google.generativeai as genai
+genai.configure(api_key="YOUR_API_KEY")
+model = genai.GenerativeModel("gemini-1.5-flash")
+response = model.generate_content("Hello")
+```
+
+**After (migrated):**
+```python
+from google import genai
+client = genai.Client(api_key="YOUR_API_KEY")
+model = genai.GenerativeModel("gemini-1.5-flash")
+response = client.models.generate_content(model="gemini-1.5-flash", contents="Hello")
+```
+
+**Supported transforms:**
+
+| Legacy pattern | New pattern |
+|---|---|
+| `import google.generativeai as genai` | `from google import genai` |
+| `genai.configure(api_key=...)` | `client = genai.Client(api_key=...)` |
+| `genai.GenerativeModel(...)` | `client.models.generate_content(model=..., ...)` |
+| `model.generate_content(...)` | `client.models.generate_content(model=..., contents=...)` |
+| `model.start_chat(...)` | `client.chats.create(model=..., config=...)` |
+| `model.count_tokens(...)` | `client.models.count_tokens(model=..., contents=...)` |
+
+**Known limitations:** `generation_config`/`safety_settings` folding into `config=` is not yet automated; `history=` in `start_chat` is stubbed.
+
 ## Imports
 
 ```python

google/genai/_migrate/cli.py

@@ -0,0 +1,99 @@
+"""CLI for google-generativeai -> google-genai migration."""
+
+import argparse
+import difflib
+import sys
+from pathlib import Path
+
+from google.genai._migrate.codemod import transform_source
+
+
+def collect_python_files(path: Path, include_tests: bool) -> list[Path]:
+    """Collect .py files from path, optionally including test_*.py files."""
+    files: list[Path] = []
+    if path.is_file() and path.suffix == ".py":
+        files.append(path)
+    elif path.is_dir():
+        for py_file in path.rglob("*.py"):
+            if include_tests or not py_file.name.startswith("test_"):
+                files.append(py_file)
+    return sorted(files)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        prog="genai-migrate",
+        description="Migrate google-generativeai code to google-genai.",
+    )
+    parser.add_argument(
+        "path",
+        type=Path,
+        help="File or directory to migrate",
+    )
+    parser.add_argument(
+        "--diff",
+        action="store_true",
+        help="Print unified diff to stdout, do not write files",
+    )
+    parser.add_argument(
+        "--in-place",
+        action="store_true",
+        help="Overwrite files in place",
+    )
+    parser.add_argument(
+        "--include-tests",
+        action="store_true",
+        help="Also process test_*.py files (default: off)",
+    )
+
+    args = parser.parse_args()
+
+    if not args.diff and not args.in_place:
+        parser.error("Must specify either --diff or --in-place")
+
+    target_path = args.path
+    if not target_path.exists():
+        parser.error(f"Path does not exist: {target_path}")
+
+    files = collect_python_files(target_path, args.include_tests)
+    if not files:
+        print("No Python files found to process.")
+        return 0
+
+    total_files = len(files)
+    changed_files = 0
+
+    for file_path in files:
+        try:
+            source = file_path.read_text(encoding="utf-8")
+        except Exception as e:
+            print(f"Error reading {file_path}: {e}", file=sys.stderr)
+            continue
+
+        new_source = transform_source(source)
+
+        if new_source == source:
+            continue
+
+        changed_files += 1
+
+        if args.diff:
+            diff = difflib.unified_diff(
+                source.splitlines(keepends=True),
+                new_source.splitlines(keepends=True),
+                fromfile=str(file_path),
+                tofile=str(file_path),
+            )
+            sys.stdout.writelines(diff)
+        elif args.in_place:
+            try:
+                file_path.write_text(new_source, encoding="utf-8")
+            except Exception as e:
+                print(f"Error writing {file_path}: {e}", file=sys.stderr)
+
+    print(f"Processed {total_files} files, {changed_files} changed.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

google/genai/_migrate/codemod.py

@@ -0,0 +1,51 @@
+# google/genai/_migrate/codemod.py
+import libcst as cst
+# from typing import Dict, Any
+
+from google.genai._migrate.visitors import (CountTokensVisitor,
+                                            GenerateContentVisitor,
+                                            GenerativeModelVisitor,
+                                            ImportAndConfigureVisitor,
+                                            StartChatVisitor)
+
+
+def transform_source(src: str) -> str:
+    """
+    Parses Python source code with libcst and applies the migration visitors.
+
+    The visitors are applied sequentially to ensure state (like the symbol_table
+    mapping variable names to model strings) is passed down correctly.
+
+    Args:
+        src: The raw Python source code string to be migrated.
+
+    Returns:
+        The migrated Python source code string, with formatting preserved.
+
+    Raises:
+        ValueError: If the source code contains invalid Python syntax.
+    """
+    if not src.strip():
+        return src
+
+    try:
+        module = cst.parse_module(src)
+    except cst.ParserSyntaxError as e:
+        raise ValueError(f"Failed to parse Python source: {e}") from e
+
+    # Shared state container: maps legacy model variable names to their string names.
+    # e.g., {"model": "'gemini-1.5-flash'"}
+    symbol_table: dict[str, str] = {}
+
+    # Apply visitors sequentially.
+    # Order matters:
+    # 1. Imports/configure must run first to establish the `client` variable.
+    # 2. GenerativeModelVisitor must run before the others to populate the symbol_table.
+    # 3. The remaining visitors use the symbol_table to rewrite method calls.
+    module = module.visit(ImportAndConfigureVisitor())
+    module = module.visit(GenerativeModelVisitor(symbol_table))
+    module = module.visit(GenerateContentVisitor(symbol_table))
+    module = module.visit(StartChatVisitor(symbol_table))
+    module = module.visit(CountTokensVisitor(symbol_table))
+
+    return module.code

google/genai/_migrate/mappings.py

@@ -0,0 +1,47 @@
+"""Legacy -> new transformation mappings for google-generativeai -> google-genai migration.
+
+Each dict represents one legacy -> new transformation with keys:
+- name: short identifier for the transformation
+- legacy_pattern: legacy code snippet (string snippet, not regex)
+- new_pattern: new code snippet (string snippet, not regex)
+- notes: additional notes about the transformation
+"""
+
+MAPPINGS: list[dict[str, str]] = [
+    {
+        "name": "import",
+        "legacy_pattern": "import google.generativeai as genai",
+        "new_pattern": "from google import genai",
+        "notes": "Change import style from module import to from google import genai",
+    },
+    {
+        "name": "configure",
+        "legacy_pattern": "genai.configure(api_key=...)",
+        "new_pattern": "client = genai.Client(api_key=...)",
+        "notes": "Replace configure() with Client instantiation",
+    },
+    {
+        "name": "model",
+        "legacy_pattern": "genai.GenerativeModel('gemini-X')",
+        "new_pattern": "client.models.generate_content(model='gemini-X', ...)",
+        "notes": "Replace GenerativeModel instantiation with client.models.generate_content",
+    },
+    {
+        "name": "generate_content",
+        "legacy_pattern": "model.generate_content(text)",
+        "new_pattern": "client.models.generate_content(model=..., contents=text)",
+        "notes": "Replace model.generate_content with client.models.generate_content",
+    },
+    {
+        "name": "start_chat",
+        "legacy_pattern": "model.start_chat(history=...)",
+        "new_pattern": "client.chats.create(model=..., config=...)",
+        "notes": "Replace model.start_chat with client.chats.create",
+    },
+    {
+        "name": "count_tokens",
+        "legacy_pattern": "model.count_tokens(text)",
+        "new_pattern": "client.models.count_tokens(model=..., contents=text)",
+        "notes": "Replace model.count_tokens with client.models.count_tokens",
+    },
+]