flyteplugins-codegen 2.1.9__py3-none-any.whl → 2.2.1__py3-none-any.whl

flyteplugins/codegen/auto_coder_agent.py

@@ -12,7 +12,7 @@ import flyte
 import litellm
 import pandas as pd
 from flyte.errors import InvalidPackageError
-from flyte.io import File
+from flyte.io import Dir, File
 from flyte.sandbox import ImageConfig
 from flyte.syncify import syncify
 
@@ -150,9 +150,9 @@ class AutoCoderAgent:
             inputs: Optional dict declaring non-sample CLI argument types
                 (e.g., `{"threshold": float, "mode": str}`).
                 Sample entries are automatically added as File inputs — don't redeclare them here.
-                Supported types: str, int, float, bool, File.
+                Supported types: str, int, float, bool, File, Dir.
             outputs: Optional dict defining output types (e.g., `{"result": str, "report": File}`).
-                Supported types: str, int, float, bool, datetime, timedelta, File.
+                Supported types: str, int, float, bool, datetime, timedelta, File, Dir.
 
         Returns:
             CodeGenEvalResult with solution and execution details.
@@ -161,7 +161,7 @@ class AutoCoderAgent:
 
         # Input validation
         if inputs:
-            supported_input_types = (str, int, float, bool, File)
+            supported_input_types = (str, int, float, bool, File, Dir)
             for input_key, input_type in inputs.items():
                 if input_type not in supported_input_types:
                     supported_names = [t.__name__ for t in supported_input_types]
@@ -238,6 +238,7 @@ class AutoCoderAgent:
                 datetime.datetime,
                 datetime.timedelta,
                 File,
+                Dir,
             )
             for output_key, output_type in outputs.items():
                 if output_type not in supported_types:

flyteplugins/codegen/generation/prompts.py

@@ -2,6 +2,8 @@
 
 from typing import Optional
 
+from flyte.io import Dir, File
+
 # Language-specific file extensions
 FILE_EXTENSIONS = {"python": ".py"}
 
@@ -70,6 +72,9 @@ def build_enhanced_prompt(
     if schema:
         enhanced_prompt += f"\n\nSchema:\n``\n{schema}\n``"
 
+    def _is_path_input_type(param_type: type) -> bool:
+        return param_type in (File, Dir) or any(name in str(param_type) for name in ("File", "Dir"))
+
     # Always add script requirement first, then user constraints
     script_constraint = (
         "REQUIRED: Your code will be saved as solution.py and imported by tests via "
@@ -85,19 +90,20 @@ def build_enhanced_prompt(
         args_list = []
         for name, param_type in inputs.items():
             type_name = param_type.__name__ if hasattr(param_type, "__name__") else str(param_type)
-            # Clarify that File inputs are received as string paths
-            if "File" in type_name:
+            # Clarify that File/Dir inputs are received as string paths
+            if _is_path_input_type(param_type):
                 args_list.append(f"--{name} (str): path to {type_name.lower()}")
             else:
                 args_list.append(f"--{name} ({type_name})")
         args_spec = ", ".join(args_list)
         script_constraint += f"Accept these command line arguments: {args_spec}. "
 
-        # Add explicit instruction for File handling
-        has_file_inputs = any("File" in str(t) for t in inputs.values())
-        if has_file_inputs:
+        # Add explicit instruction for File/Dir handling
+        has_path_inputs = any(_is_path_input_type(t) for t in inputs.values())
+        if has_path_inputs:
             script_constraint += (
-                "File arguments are string paths - use them directly with open() or other file operations."
+                "File and Dir arguments are string paths - use them directly with "
+                "open(), pathlib, or other file operations."
             )
     elif data_context:
         script_constraint += "Accept appropriate command line arguments to process the data samples."
@@ -111,16 +117,20 @@ def build_enhanced_prompt(
         output_parts = []
         for name, output_type in outputs.items():
             type_name = output_type.__name__ if hasattr(output_type, "__name__") else str(output_type)
-            if "File" in type_name:
+            if output_type is Dir or "Dir" in type_name:
+                output_parts.append(f"- {name}: create the output directory directly at /var/outputs/{name}")
+            elif _is_path_input_type(output_type):
                 output_parts.append(f"- {name}: write the output file directly to /var/outputs/{name}")
             else:
                 output_parts.append(f"- {name} ({type_name}): write the value to /var/outputs/{name}")
         output_list = "\n".join(output_parts)
         output_constraint = f"""OUTPUT REQUIREMENTS — you MUST write each output as a file under /var/outputs/:
 {output_list}
-Use this exact pattern for each output:
+Use this exact pattern for scalar outputs:
   with open('/var/outputs/<name>', 'w') as f:
       f.write(str(value))
+For File outputs, write the file directly to /var/outputs/<name>.
+For Dir outputs, create the directory directly at /var/outputs/<name>.
 /var/outputs/ already exists. NEVER delete, recreate, or modify the directory itself. Only write files into it.
 Outputs MUST be written before the script exits — do NOT just print() values."""
         all_constraints.append(output_constraint)

{flyteplugins_codegen-2.1.9.dist-info → flyteplugins_codegen-2.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyteplugins-codegen
-Version: 2.1.9
+Version: 2.2.1
 Summary: LLM-powered code generation and evaluation plugin for Flyte
 Author-email: Samhita Alla <samhita@union.ai>
 Requires-Python: >=3.10
@@ -188,8 +188,8 @@ result = await agent.generate.aio(prompt="...")
 | `schema`      | `str`                             | `None`   | Free-form context about data formats, structures, or schemas. Included verbatim in the LLM prompt.             |
 | `constraints` | `list[str]`                       | `None`   | Natural-language constraints (e.g., `"quantity must be positive"`)                                             |
 | `samples`     | `dict[str, File \| pd.DataFrame]` | `None`   | Sample data. Sampled for LLM context, converted to File inputs for the sandbox. Used as defaults at runtime.   |
-| `inputs`      | `dict[str, type]`                 | `None`   | Non-sample CLI argument types (e.g., `{"threshold": float}`). Sample entries are auto-added as File inputs.    |
-| `outputs`     | `dict[str, type]`                 | `None`   | Output types. Supported: `str, int, float, bool, datetime, timedelta, File`.                                   |
+| `inputs`      | `dict[str, type]`                 | `None`   | Non-sample CLI argument types (e.g., `{"threshold": float}`). Sample entries are auto-added as File inputs. Supported: `str, int, float, bool, File, Dir`. |
+| `outputs`     | `dict[str, type]`                 | `None`   | Output types. Supported: `str, int, float, bool, datetime, timedelta, File, Dir`.                             |
 
 ### `CodeGenEvalResult`
 

{flyteplugins_codegen-2.1.9.dist-info → flyteplugins_codegen-2.2.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 flyteplugins/codegen/__init__.py,sha256=jumoM0Po0Tx1KqatcdNi_Mk1MYCLcccnyjn6d1zurLQ,366
-flyteplugins/codegen/auto_coder_agent.py,sha256=-HkI3nM3T0Ydj8ezDprDI8DLetgQeoixBDGJXnJfmUA,46241
+flyteplugins/codegen/auto_coder_agent.py,sha256=G6F4OCJXKU-6Y16hNdodu8FtOVrOsdt2hI-lNHF40Ac,46282
 flyteplugins/codegen/core/__init__.py,sha256=YjAN0PpvkhERFmlkEq78O2q92bmgyGAVttYXeDH_hyc,355
 flyteplugins/codegen/core/types.py,sha256=kxAoKMu9tOTRNPBfiFC-dxMQOmeEawJAVRNKj42_3js,12274
 flyteplugins/codegen/data/__init__.py,sha256=UDkKuvyBEZmsPaTR6yEexcQjLutMRxuV5-pWflGc_sI,670
@@ -10,8 +10,8 @@ flyteplugins/codegen/execution/agent.py,sha256=eCIsQ09YU3OYRci1dZb8lRqxTPr21DHvW
 flyteplugins/codegen/execution/docker.py,sha256=LSmg9g7MCiD-yGAMuQT0YCh_npma_dellSh0WIvJHak,7294
 flyteplugins/codegen/generation/__init__.py,sha256=GkLiXfJeVQmLlHf4R08qzgUa5wYG97KseGjRtLCFZhU,1065
 flyteplugins/codegen/generation/llm.py,sha256=FPBt1MlW_OjwtEEcnJhE62sYrbSK4tRB3UEk80cWvgQ,44229
-flyteplugins/codegen/generation/prompts.py,sha256=Te9nK9TldLyHM2KB24PhTiWNqoEXJlW9l8nt09FbWa8,5770
-flyteplugins_codegen-2.1.9.dist-info/METADATA,sha256=k-vBB8EDdUhyslEeLYTo7ubuAoYqwRf8JHNg8ff7M80,19203
-flyteplugins_codegen-2.1.9.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-flyteplugins_codegen-2.1.9.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
-flyteplugins_codegen-2.1.9.dist-info/RECORD,,
+flyteplugins/codegen/generation/prompts.py,sha256=-Qt8M_6bkS4Jwm-bCQnCGR-oKL_mWhryCc9_OUnvXPs,6344
+flyteplugins_codegen-2.2.1.dist-info/METADATA,sha256=4A0PZUEyNYH0uBl1JLDaCtXrNz3FuDLqTas4_5o8S0A,19246
+flyteplugins_codegen-2.2.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+flyteplugins_codegen-2.2.1.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_codegen-2.2.1.dist-info/RECORD,,