flyteplugins-codegen 2.0.6__py3-none-any.whl

flyteplugins/codegen/core/__init__.py

@@ -0,0 +1,19 @@
+"""Core type definitions for LLM code generation."""
+
+from flyteplugins.codegen.core.types import (
+    CodeGenEvalResult,
+    CodePlan,
+    CodeSolution,
+    ErrorDiagnosis,
+    FixVerification,
+    TestFailure,
+)
+
+__all__ = [
+    "CodeGenEvalResult",
+    "CodePlan",
+    "CodeSolution",
+    "ErrorDiagnosis",
+    "FixVerification",
+    "TestFailure",
+]

flyteplugins/codegen/core/types.py

@@ -0,0 +1,337 @@
+from typing import Any, Literal, Optional
+
+import flyte
+from flyte.io import File
+from flyte.syncify import syncify
+from pydantic import BaseModel, Field, field_validator
+
+
+class CodePlan(BaseModel):
+    """Structured plan for the code solution."""
+
+    description: str = Field(description="Overall description of the solution")
+    approach: str = Field(description="High-level approach and algorithm to solve the problem")
+
+
+class CodeSolution(BaseModel):
+    """Structured code solution."""
+
+    language: str = Field(
+        default="python",
+        description="Programming language",
+    )
+    code: str = Field(
+        default="",
+        description="Complete executable code including imports and dependencies",
+    )
+    system_packages: list[str] = Field(
+        default_factory=list,
+        description="System packages needed (e.g., gcc, build-essential, curl)",
+    )
+
+    @field_validator("language", mode="before")
+    @classmethod
+    def normalize_language(cls, v: str) -> str:
+        return v.strip().lower()
+
+
+class CodeGenEvalResult(BaseModel):
+    """Result from code generation and evaluation."""
+
+    plan: Optional[CodePlan] = None
+    solution: CodeSolution
+    tests: Optional[str] = None
+    success: bool
+    output: str
+    exit_code: int
+    error: Optional[str] = None
+    attempts: int = 1
+    conversation_history: list[dict[str, str]] = Field(default_factory=list)
+    detected_packages: list[str] = Field(
+        default_factory=list,
+        description="Language packages detected by LLM from imports",
+    )
+    detected_system_packages: list[str] = Field(default_factory=list, description="System packages detected by LLM")
+    image: Optional[str] = Field(
+        default=None,
+        description="The Flyte Image built with all dependencies",
+    )
+    total_input_tokens: int = Field(
+        default=0,
+        description="Total input tokens used across all LLM calls",
+    )
+    total_output_tokens: int = Field(
+        default=0,
+        description="Total output tokens used across all LLM calls",
+    )
+    declared_inputs: Optional[dict[str, type]] = Field(
+        default=None,
+        description="Input types (user-provided or inferred from samples)",
+    )
+    declared_outputs: Optional[dict[str, type]] = Field(
+        default=None,
+        description="Output types declared by user",
+    )
+    data_context: Optional[str] = Field(
+        default=None,
+        description="Extracted data context (schema, stats, patterns, samples) used for code generation",
+    )
+    original_samples: Optional[dict[str, File]] = Field(
+        default=None,
+        description="Sample data converted to Files (defaults for run()/as_task())",
+    )
+    generated_schemas: Optional[dict[str, str]] = Field(
+        default=None,
+        description="Auto-generated Pandera schemas (as Python code strings) for validating data inputs",
+    )
+
+    def as_task(
+        self,
+        name: str = "run_code_on_real_data",
+        resources: Optional[flyte.Resources] = None,
+        retries: int = 0,
+        timeout: Optional[int] = None,
+        env_vars: Optional[dict[str, str]] = None,
+        secrets: Optional[list] = None,
+        cache: str = "auto",
+    ):
+        """Create a sandbox that runs the generated code in an isolated sandbox.
+
+        The generated code will write outputs to /var/outputs/{output_name} files.
+        Returns a callable wrapper that automatically provides the script file.
+
+        Args:
+            name: Name for the sandbox
+            resources: Optional resources for the task
+            retries: Number of retries for the task. Defaults to 0.
+            timeout: Timeout in seconds. Defaults to None.
+            env_vars: Environment variables to pass to the sandbox.
+            secrets: flyte.Secret objects to make available.
+            cache: CacheRequest: "auto", "override", or "disable". Defaults to "auto".
+
+        Returns:
+            Callable task wrapper with the default inputs baked in. Call with your other declared inputs.
+        """
+        if not self.success:
+            raise ValueError("Cannot create task from failed code generation")
+
+        if not self.image:
+            raise ValueError("No image available - code generation did not build an image")
+
+        sandbox = flyte.sandbox.create(
+            name=name,
+            code=self.solution.code,
+            inputs=self.declared_inputs or {},
+            outputs=self.declared_outputs or {},
+            auto_io=False,
+            resources=resources or flyte.Resources(cpu=1, memory="1Gi"),
+            retries=retries,
+            timeout=timeout,
+            env_vars=env_vars,
+            secrets=secrets,
+            cache=cache,
+        )
+
+        image = self.image
+
+        # If we have samples, wrap to inject sample values as defaults
+        if self.original_samples:
+            sample_defaults = dict(self.original_samples)
+
+            @syncify
+            async def task_with_defaults(**kwargs):
+                merged = {**sample_defaults, **kwargs}
+                return await sandbox.run.aio(image=image, **merged)
+
+            return task_with_defaults
+
+        @syncify
+        async def task(**kwargs):
+            return await sandbox.run.aio(image=image, **kwargs)
+
+        return task
+
+    async def run(
+        self,
+        *,
+        name: str = "run_code_on_real_data",
+        resources: Optional[flyte.Resources] = None,
+        retries: int = 0,
+        timeout: Optional[int] = None,
+        env_vars: Optional[dict[str, str]] = None,
+        secrets: Optional[list] = None,
+        cache: str = "auto",
+        **overrides,
+    ) -> Any:
+        """Run generated code in an isolated sandbox (one-off execution).
+
+        If samples were provided during generate(), they are used as defaults.
+        Override any input by passing it as a keyword argument. If no samples
+        exist, all declared inputs must be provided via ``**overrides``.
+
+        Args:
+            name: Name for the sandbox
+            resources: Optional resources for the task
+            retries: Number of retries for the task. Defaults to 0.
+            timeout: Timeout in seconds. Defaults to None.
+            env_vars: Environment variables to pass to the sandbox.
+            secrets: flyte.Secret objects to make available.
+            cache: CacheRequest: "auto", "override", or "disable". Defaults to "auto".
+            **overrides: Input values. Merged on top of sample defaults (if any).
+
+        Returns:
+            Tuple of typed outputs.
+        """
+        if not self.success:
+            raise ValueError("Cannot run failed code generation")
+
+        if not self.image:
+            raise ValueError("No image available - code generation did not build an image")
+
+        sandbox = flyte.sandbox.create(
+            name=name,
+            code=self.solution.code,
+            inputs=self.declared_inputs or {},
+            outputs=self.declared_outputs or {},
+            auto_io=False,
+            resources=resources or flyte.Resources(cpu=1, memory="1Gi"),
+            retries=retries,
+            timeout=timeout,
+            env_vars=env_vars,
+            secrets=secrets,
+            cache=cache,
+        )
+
+        run_data = {**(self.original_samples or {}), **overrides}
+        return await sandbox.run.aio(image=self.image, **run_data)
+
+
+# Apply syncify after class definition to avoid Pydantic field detection
+CodeGenEvalResult.run = syncify(CodeGenEvalResult.run)
+
+
+class TestFailure(BaseModel):
+    """Individual test failure with diagnosis."""
+
+    test_name: str = Field(description="Name of the failing test")
+    error_message: str = Field(
+        description="The exact final error message from test output "
+        "(e.g., 'RecursionError: maximum recursion depth exceeded')"
+    )
+    expected_behavior: str = Field(description="What this test expected to happen")
+    actual_behavior: str = Field(description="What actually happened when the code ran")
+    root_cause: str = Field(description="Why the test failed (quote the exact code that's wrong)")
+    suggested_fix: str = Field(description="Specific code changes using format: Replace `current code` with `new code`")
+    error_type: Literal["environment", "logic", "test_error"] = Field(
+        description="Type of error: 'environment' (missing packages/dependencies), "
+        "'logic' (bug in solution code), or 'test_error' (bug in test code)"
+    )
+
+
+class ErrorDiagnosis(BaseModel):
+    """Structured diagnosis of execution errors."""
+
+    failures: list[TestFailure] = Field(description="Individual test failures with their diagnoses")
+    needs_system_packages: list[str] = Field(
+        default_factory=list,
+        description="System packages needed (e.g., gcc, pkg-config).",
+    )
+    needs_language_packages: list[str] = Field(
+        default_factory=list,
+        description="Language packages needed.",
+    )
+    needs_additional_commands: list[str] = Field(
+        default_factory=list,
+        description="Additional RUN commands (e.g., apt-get update, mkdir /data, wget files).",
+    )
+
+
+class FixVerification(BaseModel):
+    """Verification that fixes were applied to code."""
+
+    all_fixes_applied: bool = Field(description="True if all suggested fixes are present in the new code")
+    applied_fixes: list[str] = Field(
+        default_factory=list,
+        description="List of fixes that were successfully applied (by test name)",
+    )
+    missing_fixes: list[str] = Field(
+        default_factory=list,
+        description="List of fixes that are still missing (by test name)",
+    )
+    explanation: str = Field(description="Brief explanation of what was checked and what's missing (if anything)")
+
+
+class TestFunctionPatch(BaseModel):
+    """A single fixed test function."""
+
+    test_name: str = Field(description="Name of the test function (e.g. test_basic_analysis)")
+    fixed_code: str = Field(description="Complete fixed function body including the def line and decorators")
+
+
+class TestFixResponse(BaseModel):
+    """Response containing only the fixed test functions."""
+
+    patches: list[TestFunctionPatch] = Field(description="List of fixed test functions")
+
+
+class _PackageReplacementResponse(BaseModel):
+    """Response format for suggesting a replacement system package."""
+
+    replacement: Optional[str] = Field(
+        default=None,
+        description="Correct Debian/Ubuntu apt package name, or null if no system package is needed",
+    )
+
+
+class _PackageDetectionResponse(BaseModel):
+    """Response format for LLM package detection."""
+
+    packages: list[str] = Field(
+        default_factory=list,
+        description="List of third-party package names",
+    )
+
+
+class _TestCodeResponse(BaseModel):
+    """Response format for LLM test generation."""
+
+    test_code: str = Field(description="Complete test code")
+
+
+class _ConstraintParameters(BaseModel):
+    """Parameters for a constraint check. Only the fields relevant to the check_type should be set."""
+
+    value: Optional[float] = Field(
+        default=None,
+        description="Threshold value for greater_than or less_than checks",
+    )
+    min: Optional[float] = Field(
+        default=None,
+        description="Minimum value for between checks",
+    )
+    max: Optional[float] = Field(
+        default=None,
+        description="Maximum value for between checks",
+    )
+    pattern: Optional[str] = Field(
+        default=None,
+        description="Regex pattern for regex checks",
+    )
+    values: Optional[list[str]] = Field(
+        default=None,
+        description="Allowed values for isin checks",
+    )
+
+
+class _ConstraintParse(BaseModel):
+    """LLM response for parsing a constraint into Pandera check."""
+
+    column_name: str = Field(description="Name of the column this constraint applies to")
+    check_type: Literal["greater_than", "less_than", "between", "regex", "isin", "not_null", "none"] = Field(
+        description="Type of check to apply"
+    )
+    parameters: _ConstraintParameters = Field(
+        default_factory=_ConstraintParameters,
+        description="Parameters for the check. Set only the fields relevant to the check_type.",
+    )
+    explanation: str = Field(description="Brief explanation of what check will be applied")

flyteplugins/codegen/data/__init__.py

@@ -0,0 +1,27 @@
+"""Data extraction and schema inference."""
+
+from flyteplugins.codegen.data.extraction import (
+    extract_data_context,
+    extract_dataframe_context,
+    extract_file_context,
+    is_dataframe,
+)
+from flyteplugins.codegen.data.schema import (
+    apply_parsed_constraint,
+    apply_user_constraints,
+    extract_token_usage,
+    infer_conservative_schema,
+    parse_constraint_with_llm,
+)
+
+__all__ = [
+    "apply_parsed_constraint",
+    "apply_user_constraints",
+    "extract_data_context",
+    "extract_dataframe_context",
+    "extract_file_context",
+    "extract_token_usage",
+    "infer_conservative_schema",
+    "is_dataframe",
+    "parse_constraint_with_llm",
+]

flyteplugins/codegen/data/extraction.py

@@ -0,0 +1,281 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+import aiofiles
+import flyte
+import pandas as pd
+import pandera.pandas as pa
+from flyte.io import File
+
+from flyteplugins.codegen.data.schema import (
+    apply_user_constraints,
+    infer_conservative_schema,
+    schema_to_script,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def is_dataframe(obj) -> bool:
+    """Check if object is a pandas DataFrame.
+
+    Args:
+        obj: Object to check
+
+    Returns:
+        True if object is a DataFrame
+    """
+    try:
+        return isinstance(obj, pd.DataFrame)
+    except ImportError:
+        return False
+
+
+async def extract_dataframe_context(
+    df, name: str, max_sample_rows: int = 5, schema: Optional[pa.DataFrameSchema] = None
+) -> str:
+    """Extract comprehensive context from DataFrame.
+
+    Args:
+        df: pandas DataFrame
+        name: Name of the data input
+        max_sample_rows: Number of sample rows to include
+        schema: Optional Pandera schema to include in context
+
+    Returns:
+        Formatted string with all extracted context
+    """
+    context_parts = []
+
+    # 1. Structural Context
+    context_parts.append(f"## Data: {name}")
+    context_parts.append(f"Shape: {df.shape[0]:,} rows x {df.shape[1]} columns")
+
+    # Include Pandera schema if provided (use Pandera's built-in formatter)
+    if schema:
+        context_parts.append(f"\nPandera Schema for {name} (use for validation):")
+        context_parts.append("```python")
+        context_parts.append(schema_to_script(schema))
+        context_parts.append("```")
+
+    # 2. Statistical Context
+    context_parts.append("\nStatistical Summary:")
+
+    # Numeric columns
+    numeric_cols = df.select_dtypes(include=["number"]).columns
+    if len(numeric_cols) > 0:
+        context_parts.append("  Numeric columns:")
+        desc = df[numeric_cols].describe()
+        for col in numeric_cols:
+            stats = desc[col]
+            context_parts.append(
+                f"    {col}: min={stats['min']:.2g}, max={stats['max']:.2g}, "
+                f"mean={stats['mean']:.2g}, median={stats['50%']:.2g}"
+            )
+
+    # Categorical/Object columns
+    cat_cols = df.select_dtypes(include=["object", "category"]).columns
+    if len(cat_cols) > 0:
+        context_parts.append("  Categorical columns:")
+        for col in cat_cols:
+            unique_count = df[col].nunique()
+            total_count = len(df[col].dropna())
+            if unique_count <= 20 and total_count > 0:
+                # Show value counts for low-cardinality columns
+                top_values = df[col].value_counts().head(5)
+                top_str = ", ".join([f"'{k}': {v}" for k, v in top_values.items()])
+                context_parts.append(f"    {col}: {unique_count} unique values. Top 5: {{{top_str}}}")
+            else:
+                context_parts.append(f"    {col}: {unique_count} unique values")
+
+    # DateTime columns
+    date_cols = df.select_dtypes(include=["datetime64"]).columns
+    if len(date_cols) > 0:
+        context_parts.append("  DateTime columns:")
+        for col in date_cols:
+            min_date = df[col].min()
+            max_date = df[col].max()
+            context_parts.append(f"    {col}: {min_date} to {max_date}")
+
+    # 3. Behavioral Context (patterns, invariants)
+    context_parts.append("\nData Patterns:")
+
+    # Check for duplicates
+    dup_count = df.duplicated().sum()
+    if dup_count > 0:
+        context_parts.append(f"  - {dup_count:,} duplicate rows ({dup_count / len(df) * 100:.1f}%)")
+
+    # Check for potential ID columns
+    for col in df.columns:
+        if df[col].nunique() == len(df) and not df[col].isna().any():
+            context_parts.append(f"  - '{col}' appears to be a unique identifier")
+            break
+
+    # 4. Representative Samples
+    context_parts.append(f"\nRepresentative Samples ({max_sample_rows} rows):")
+
+    # Sample strategy: first few + random + edge cases
+    sample_indices = []
+
+    # First rows
+    sample_indices.extend(range(min(2, len(df))))
+
+    # Random sample
+    if len(df) > max_sample_rows:
+        remaining = max_sample_rows - len(sample_indices)
+        random_indices = df.sample(n=remaining).index.tolist()
+        sample_indices.extend(random_indices)
+    else:
+        sample_indices = list(range(len(df)))
+
+    sample_df = df.iloc[sample_indices[:max_sample_rows]]
+
+    # Format as CSV
+    context_parts.append(sample_df.to_csv(index=False))
+
+    return "\n".join(context_parts)
+
+
+async def extract_file_context(file: File, name: str, max_sample_rows: int = 5) -> str:
+    """Extract context from non-tabular files (text, binary, unknown formats).
+
+    This is a fallback for files that can't be loaded as DataFrames.
+    Structured files (CSV, Parquet, JSON, Excel) are handled by extract_data_context()
+    with Pandera schema inference.
+
+    Args:
+        file: File to extract context from
+        name: Name of the data input
+        max_sample_rows: Number of sample rows to include
+
+    Returns:
+        Formatted string with all extracted context
+    """
+    local_path = await file.download()
+    file_ext = Path(local_path).suffix.lower()
+
+    # Try to read as text file
+    try:
+        async with aiofiles.open(local_path, "r", encoding="utf-8", errors="ignore") as f:
+            lines = []
+            for _ in range(max_sample_rows):
+                line = await f.readline()
+                if not line:
+                    break
+                lines.append(line)
+
+        context_parts = [
+            f"## Data: {name}",
+            f"Type: Text file ({file_ext})",
+            f"Lines: {len(lines)}",
+            f"\nFirst {max_sample_rows} lines:",
+            "".join(lines),
+        ]
+        return "\n".join(context_parts)
+
+    except Exception:
+        # Binary or unreadable file
+        file_size = Path(local_path).stat().st_size  # noqa: ASYNC240
+        context_parts = [
+            f"## Data: {name}",
+            f"Type: Binary/Unknown ({file_ext})",
+            f"Size: {file_size:,} bytes",
+            "\n(Unable to extract text preview)",
+        ]
+        return "\n".join(context_parts)
+
+
+@flyte.trace
+async def extract_data_context(
+    data: dict[str, pd.DataFrame | File],
+    max_sample_rows: int = 5,
+    constraints: Optional[list[str]] = None,
+    model: Optional[str] = None,
+    litellm_params: Optional[dict] = None,
+) -> tuple[str, dict[str, str], int, int]:
+    """Extract comprehensive context from data inputs with Pandera schema inference.
+
+    Extracts:
+    1. Structural context (schema, types, shape)
+    2. Statistical context (distributions, ranges)
+    3. Behavioral context (patterns, invariants)
+    4. Operational context (scale, nulls)
+    5. Representative samples
+    6. Pandera schemas (inference + user constraints), returned as Python code strings
+
+    Args:
+        data: Dict of data inputs (File or DataFrame)
+        max_sample_rows: Number of sample rows to include
+        constraints: Optional list of user constraints to apply to schemas
+        model: LLM model for constraint parsing (required if constraints provided)
+        litellm_params: Optional LiteLLM parameters
+
+    Returns:
+        Tuple of (context_string, schemas_as_code_dict, total_input_tokens, total_output_tokens)
+    """
+    context_parts = []
+    schemas: dict[str, str] = {}
+    total_input_tokens = 0
+    total_output_tokens = 0
+
+    for name, value in data.items():
+        df = None
+
+        if isinstance(value, File):
+            # Load file as DataFrame for schema inference
+            local_path = await value.download()
+            file_ext = Path(local_path).suffix.lower()
+
+            try:
+                if file_ext in [".csv", ".tsv"]:
+                    delimiter = "\t" if file_ext == ".tsv" else ","
+                    df = pd.read_csv(local_path, delimiter=delimiter, nrows=10000)
+                elif file_ext in [".parquet", ".pq"]:
+                    df = pd.read_parquet(local_path)
+                    if len(df) > 10000:
+                        df = df.sample(n=10000)
+                elif file_ext == ".json":
+                    try:
+                        df = pd.read_json(local_path, lines=True, nrows=10000)
+                    except Exception:
+                        df = pd.read_json(local_path)
+                elif file_ext in [".xlsx", ".xls"]:
+                    df = pd.read_excel(local_path, nrows=10000)
+                else:
+                    # Non-tabular file (e.g., .log, .txt) — extract text context
+                    context = await extract_file_context(value, name, max_sample_rows)
+                    context_parts.append(context)
+                    continue
+            except Exception as e:
+                logger.warning(f"Failed to load {name} as DataFrame for schema inference: {e}")
+                # Fall back to non-schema extraction
+                context = await extract_file_context(value, name, max_sample_rows)
+                context_parts.append(context)
+                continue
+
+        elif is_dataframe(value):
+            df = value
+        else:
+            context_parts.append(f"## Data: {name}\nType: {type(value)}\n(Unsupported type)")
+            continue
+
+        if df is not None:
+            # Infer Pandera schema
+            schema = infer_conservative_schema(df)
+
+            # Apply user constraints if provided
+            if constraints and model:
+                schema, in_tok, out_tok = await apply_user_constraints(schema, constraints, name, model, litellm_params)
+                total_input_tokens += in_tok
+                total_output_tokens += out_tok
+
+            # Convert to code string for serialization
+            schemas[name] = schema_to_script(schema)
+
+            # Extract context with schema
+            context = await extract_dataframe_context(df, name, max_sample_rows, schema)
+            context_parts.append(context)
+
+    context_str = "\n\n" + "=" * 80 + "\n\n".join(context_parts)
+    return context_str, schemas, total_input_tokens, total_output_tokens