codeshift 0.2.0__py3-none-any.whl

codeshift/migrator/llm_migrator.py

@@ -0,0 +1,320 @@
+"""LLM-based migration for complex cases.
+
+LLM migrations (Tier 2/3) are routed through the PyResolve API,
+which handles authentication, quota checking, and billing.
+Users must have a Pro or Unlimited subscription to use LLM features.
+
+Tier 1 (deterministic AST transforms) remains free and local.
+"""
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+from codeshift.migrator.ast_transforms import (
+    TransformChange,
+    TransformResult,
+    TransformStatus,
+)
+from codeshift.utils.api_client import PyResolveAPIClient, get_api_client
+from codeshift.utils.cache import LLMCache, get_llm_cache
+from codeshift.validator.syntax_checker import quick_syntax_check
+
+
+@dataclass
+class LLMMigrationResult:
+    """Result of an LLM-based migration."""
+
+    original_code: str
+    migrated_code: str
+    success: bool
+    used_cache: bool = False
+    error: str | None = None
+    validation_passed: bool = True
+    usage: dict | None = None
+
+
+class LLMMigrator:
+    """Handles complex migrations using LLM via the PyResolve API.
+
+    LLM migrations require authentication and a Pro or Unlimited subscription.
+    All LLM calls are routed through the PyResolve API, which handles:
+    - Authentication and authorization
+    - Quota checking and billing
+    - Server-side Anthropic API calls
+
+    This ensures users cannot bypass the subscription model by using
+    their own API keys.
+    """
+
+    def __init__(
+        self,
+        client: PyResolveAPIClient | None = None,
+        cache: LLMCache | None = None,
+        use_cache: bool = True,
+        validate_output: bool = True,
+    ):
+        """Initialize the LLM migrator.
+
+        Args:
+            client: API client to use. Defaults to singleton.
+            cache: Cache to use. Defaults to singleton.
+            use_cache: Whether to use caching
+            validate_output: Whether to validate migrated code syntax
+        """
+        self.client = client or get_api_client()
+        self.cache = cache or get_llm_cache() if use_cache else None
+        self.use_cache = use_cache
+        self.validate_output = validate_output
+
+    @property
+    def is_available(self) -> bool:
+        """Check if LLM migration is available (user is authenticated)."""
+        return self.client.is_available
+
+    def migrate(
+        self,
+        code: str,
+        library: str,
+        from_version: str,
+        to_version: str,
+        context: str | None = None,
+    ) -> LLMMigrationResult:
+        """Migrate code using the LLM via PyResolve API.
+
+        Args:
+            code: Source code to migrate
+            library: Library being upgraded
+            from_version: Current version
+            to_version: Target version
+            context: Optional context about the code
+
+        Returns:
+            LLMMigrationResult with the migrated code
+        """
+        if not self.is_available:
+            return LLMMigrationResult(
+                original_code=code,
+                migrated_code=code,
+                success=False,
+                error=(
+                    "LLM migrations require authentication. "
+                    "Run 'codeshift login' to authenticate, then upgrade to Pro tier."
+                ),
+            )
+
+        # Check cache first
+        if self.use_cache and self.cache:
+            cached = self.cache.get_migration(code, library, from_version, to_version)
+            if cached:
+                return LLMMigrationResult(
+                    original_code=code,
+                    migrated_code=cached,
+                    success=True,
+                    used_cache=True,
+                )
+
+        # Call PyResolve API (which calls Anthropic server-side)
+        response = self.client.migrate_code(
+            code=code,
+            library=library,
+            from_version=from_version,
+            to_version=to_version,
+            context=context,
+        )
+
+        if not response.success:
+            return LLMMigrationResult(
+                original_code=code,
+                migrated_code=code,
+                success=False,
+                error=response.error,
+            )
+
+        migrated_code = response.content
+
+        # Validate syntax
+        validation_passed = True
+        if self.validate_output:
+            if not quick_syntax_check(migrated_code):
+                validation_passed = False
+                # Try to fix common issues
+                migrated_code = self._attempt_fix(migrated_code)
+                if not quick_syntax_check(migrated_code):
+                    return LLMMigrationResult(
+                        original_code=code,
+                        migrated_code=code,
+                        success=False,
+                        error="LLM output has syntax errors",
+                        validation_passed=False,
+                    )
+
+        # Cache the result
+        if self.use_cache and self.cache:
+            self.cache.set_migration(code, library, from_version, to_version, migrated_code)
+
+        return LLMMigrationResult(
+            original_code=code,
+            migrated_code=migrated_code,
+            success=True,
+            validation_passed=validation_passed,
+            usage=response.usage,
+        )
+
+    def _attempt_fix(self, code: str) -> str:
+        """Attempt to fix common syntax issues in LLM output.
+
+        Args:
+            code: Code with potential issues
+
+        Returns:
+            Fixed code (or original if unfixable)
+        """
+        # Common fixes
+        fixes = [
+            # Remove trailing incomplete lines
+            (r"\n\s*$", "\n"),
+            # Fix unclosed strings (simple cases)
+            (r'(["\'])([^"\'\n]*?)$', r"\1\2\1"),
+        ]
+
+        fixed = code
+        for pattern, replacement in fixes:
+            fixed = re.sub(pattern, replacement, fixed)
+
+        return fixed
+
+    def explain_migration(
+        self,
+        original: str,
+        transformed: str,
+        library: str,
+    ) -> str | None:
+        """Get an explanation of a migration change.
+
+        Args:
+            original: Original code
+            transformed: Transformed code
+            library: Library being upgraded
+
+        Returns:
+            Explanation string or None if unavailable
+        """
+        if not self.is_available:
+            return None
+
+        response = self.client.explain_change(original, transformed, library)
+        if response.success:
+            return response.content.strip()
+        return None
+
+
+def migrate_with_llm_fallback(
+    code: str,
+    library: str,
+    from_version: str,
+    to_version: str,
+    deterministic_result: TransformResult | None = None,
+) -> TransformResult:
+    """Migrate code with LLM as a fallback for failures.
+
+    IMPORTANT: LLM migrations require Pro tier or higher subscription.
+    If the user is not authenticated or doesn't have the required tier,
+    only deterministic transforms will be applied.
+
+    Args:
+        code: Source code to migrate
+        library: Library being upgraded
+        from_version: Current version
+        to_version: Target version
+        deterministic_result: Optional result from deterministic transform
+
+    Returns:
+        TransformResult combining deterministic and LLM migrations
+    """
+    # If deterministic transform succeeded fully, use it
+    if deterministic_result and deterministic_result.status == TransformStatus.SUCCESS:
+        return deterministic_result
+
+    # Try LLM migration
+    migrator = LLMMigrator()
+
+    if not migrator.is_available:
+        # Return deterministic result if available, or no changes
+        if deterministic_result:
+            # Add a note about LLM not being available
+            errors = list(deterministic_result.errors)
+            errors.append(
+                "LLM fallback not available. Run 'codeshift login' and upgrade to Pro "
+                "for LLM-powered migrations."
+            )
+            return TransformResult(
+                file_path=deterministic_result.file_path,
+                status=deterministic_result.status,
+                original_code=deterministic_result.original_code,
+                transformed_code=deterministic_result.transformed_code,
+                changes=deterministic_result.changes,
+                errors=errors,
+            )
+        return TransformResult(
+            file_path=Path("<unknown>"),
+            status=TransformStatus.NO_CHANGES,
+            original_code=code,
+            transformed_code=code,
+            errors=[
+                "LLM not available. Run 'codeshift login' and upgrade to Pro "
+                "for LLM-powered migrations."
+            ],
+        )
+
+    # Use deterministic result as base if available
+    base_code = deterministic_result.transformed_code if deterministic_result else code
+    context = None
+
+    if deterministic_result and deterministic_result.errors:
+        context = f"Deterministic transform had issues: {', '.join(deterministic_result.errors)}"
+
+    result = migrator.migrate(
+        code=base_code,
+        library=library,
+        from_version=from_version,
+        to_version=to_version,
+        context=context,
+    )
+
+    # Combine results
+    changes = []
+    if deterministic_result:
+        changes.extend(deterministic_result.changes)
+
+    if result.success and result.migrated_code != base_code:
+        changes.append(
+            TransformChange(
+                description="LLM-assisted migration",
+                line_number=1,
+                original="(various)",
+                replacement="(LLM migrated)",
+                transform_name="llm_migration",
+                confidence=0.8,  # Lower confidence for LLM
+                notes="This change was made by the LLM and should be reviewed carefully",
+            )
+        )
+
+    status = TransformStatus.SUCCESS if result.success else TransformStatus.PARTIAL
+    if not changes:
+        status = TransformStatus.NO_CHANGES
+
+    errors = []
+    if deterministic_result:
+        errors.extend(deterministic_result.errors)
+    if result.error:
+        errors.append(result.error)
+
+    return TransformResult(
+        file_path=deterministic_result.file_path if deterministic_result else Path("<unknown>"),
+        status=status,
+        original_code=code,
+        transformed_code=result.migrated_code if result.success else base_code,
+        changes=changes,
+        errors=errors,
+    )

codeshift/migrator/transforms/__init__.py

@@ -0,0 +1,19 @@
+"""Library-specific transformation modules."""
+
+from codeshift.migrator.transforms.fastapi_transformer import FastAPITransformer
+from codeshift.migrator.transforms.pandas_transformer import (
+    PandasAppendTransformer,
+    PandasTransformer,
+)
+from codeshift.migrator.transforms.pydantic_v1_to_v2 import PydanticV1ToV2Transformer
+from codeshift.migrator.transforms.requests_transformer import RequestsTransformer
+from codeshift.migrator.transforms.sqlalchemy_transformer import SQLAlchemyTransformer
+
+__all__ = [
+    "PydanticV1ToV2Transformer",
+    "FastAPITransformer",
+    "SQLAlchemyTransformer",
+    "PandasTransformer",
+    "PandasAppendTransformer",
+    "RequestsTransformer",
+]

codeshift/migrator/transforms/fastapi_transformer.py

@@ -0,0 +1,174 @@
+"""FastAPI transformation using LibCST."""
+
+import libcst as cst
+
+from codeshift.migrator.ast_transforms import BaseTransformer
+
+
+class FastAPITransformer(BaseTransformer):
+    """Transform FastAPI code for version upgrades (0.99 to 0.100+)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._in_fastapi_context = False
+
+    def leave_ImportFrom(
+        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
+    ) -> cst.ImportFrom:
+        """Transform starlette imports to fastapi imports."""
+        if original_node.module is None:
+            return updated_node
+
+        module_name = self._get_module_name(original_node.module)
+
+        # Transform starlette.responses imports
+        if module_name == "starlette.responses":
+            self.record_change(
+                description="Import from fastapi.responses instead of starlette.responses",
+                line_number=1,
+                original=f"from {module_name}",
+                replacement="from fastapi.responses",
+                transform_name="starlette_to_fastapi_responses",
+            )
+            return updated_node.with_changes(
+                module=cst.Attribute(
+                    value=cst.Name("fastapi"),
+                    attr=cst.Name("responses"),
+                )
+            )
+
+        # Transform starlette.requests imports
+        if module_name == "starlette.requests":
+            self.record_change(
+                description="Import Request from fastapi instead of starlette.requests",
+                line_number=1,
+                original=f"from {module_name}",
+                replacement="from fastapi",
+                transform_name="starlette_to_fastapi_request",
+            )
+            return updated_node.with_changes(module=cst.Name("fastapi"))
+
+        # Transform starlette.websockets imports
+        if module_name == "starlette.websockets":
+            self.record_change(
+                description="Import WebSocket from fastapi instead of starlette.websockets",
+                line_number=1,
+                original=f"from {module_name}",
+                replacement="from fastapi",
+                transform_name="starlette_to_fastapi_websocket",
+            )
+            return updated_node.with_changes(module=cst.Name("fastapi"))
+
+        # Transform starlette.status imports
+        if module_name == "starlette.status":
+            self.record_change(
+                description="Import status from fastapi instead of starlette",
+                line_number=1,
+                original=f"from {module_name}",
+                replacement="from fastapi import status",
+                transform_name="starlette_to_fastapi_status",
+            )
+            return updated_node.with_changes(module=cst.Name("fastapi"))
+
+        return updated_node
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
+        """Transform FastAPI function calls."""
+        # Handle Field, Query, Path, Body regex -> pattern
+        if isinstance(updated_node.func, cst.Name):
+            func_name = updated_node.func.value
+            if func_name in ("Field", "Query", "Path", "Body"):
+                new_args = []
+                changed = False
+                for arg in updated_node.args:
+                    if isinstance(arg.keyword, cst.Name) and arg.keyword.value == "regex":
+                        new_args.append(arg.with_changes(keyword=cst.Name("pattern")))
+                        changed = True
+                        self.record_change(
+                            description=f"Rename {func_name}(regex=...) to {func_name}(pattern=...)",
+                            line_number=1,
+                            original=f"{func_name}(regex=...)",
+                            replacement=f"{func_name}(pattern=...)",
+                            transform_name=f"{func_name.lower()}_regex_to_pattern",
+                        )
+                    else:
+                        new_args.append(arg)
+
+                if changed:
+                    return updated_node.with_changes(args=new_args)
+
+        # Handle Depends(use_cache=...) -> Depends(use_cached=...)
+        if isinstance(updated_node.func, cst.Name) and updated_node.func.value == "Depends":
+            new_args = []
+            changed = False
+            for arg in updated_node.args:
+                if isinstance(arg.keyword, cst.Name) and arg.keyword.value == "use_cache":
+                    new_args.append(arg.with_changes(keyword=cst.Name("use_cached")))
+                    changed = True
+                    self.record_change(
+                        description="Rename Depends(use_cache=...) to Depends(use_cached=...)",
+                        line_number=1,
+                        original="Depends(use_cache=...)",
+                        replacement="Depends(use_cached=...)",
+                        transform_name="depends_use_cache_rename",
+                    )
+                else:
+                    new_args.append(arg)
+
+            if changed:
+                return updated_node.with_changes(args=new_args)
+
+        # Handle FastAPI(openapi_prefix=...) -> FastAPI(root_path=...)
+        if isinstance(updated_node.func, cst.Name) and updated_node.func.value == "FastAPI":
+            new_args = []
+            changed = False
+            for arg in updated_node.args:
+                if isinstance(arg.keyword, cst.Name) and arg.keyword.value == "openapi_prefix":
+                    new_args.append(arg.with_changes(keyword=cst.Name("root_path")))
+                    changed = True
+                    self.record_change(
+                        description="Rename openapi_prefix to root_path",
+                        line_number=1,
+                        original="FastAPI(openapi_prefix=...)",
+                        replacement="FastAPI(root_path=...)",
+                        transform_name="openapi_prefix_to_root_path",
+                    )
+                else:
+                    new_args.append(arg)
+
+            if changed:
+                return updated_node.with_changes(args=new_args)
+
+        return updated_node
+
+    def _get_module_name(self, module: cst.BaseExpression) -> str:
+        """Get the full module name from a Name or Attribute node."""
+        if isinstance(module, cst.Name):
+            return str(module.value)
+        elif isinstance(module, cst.Attribute):
+            return f"{self._get_module_name(module.value)}.{module.attr.value}"
+        return ""
+
+
+def transform_fastapi(source_code: str) -> tuple[str, list]:
+    """Transform FastAPI code.
+
+    Args:
+        source_code: The source code to transform
+
+    Returns:
+        Tuple of (transformed_code, list of changes)
+    """
+    try:
+        tree = cst.parse_module(source_code)
+    except cst.ParserSyntaxError:
+        return source_code, []
+
+    transformer = FastAPITransformer()
+    transformer.set_source(source_code)
+
+    try:
+        transformed_tree = tree.visit(transformer)
+        return transformed_tree.code, transformer.changes
+    except Exception:
+        return source_code, []