Improving external SDK for function

Author: ritaagarwala-sf

src/datacustomcode/cli.py

@@ -74,6 +74,29 @@ def _configure_client_credentials(
     )
 
 
+def _generate_function_test_file(entrypoint_path: str) -> Optional[str]:
+    """Generate test.json file for a function.
+
+    Args:
+        entrypoint_path: Path to the function's entrypoint.py
+
+    Returns:
+        Path to generated test.json, or None if generation failed
+    """
+    from datacustomcode.template import generate_test_json
+
+    tests_dir = os.path.join(os.path.dirname(entrypoint_path), "tests")
+    os.makedirs(tests_dir, exist_ok=True)
+    test_json_path = os.path.join(tests_dir, "test.json")
+
+    try:
+        generate_test_json(entrypoint_path, test_json_path)
+        return test_json_path
+    except Exception as e:
+        logger.warning(f"Could not generate test.json: {e}")
+        return None
+
+
 @cli.command()
 @click.option("--profile", default="default", help="Credential profile name")
 @click.option(
@@ -162,7 +185,6 @@ def zip(path: str, network: str):
 
     Choose based on your workload requirements.""",
 )
-@click.option("--function-invoke-opt")
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -176,13 +198,14 @@ def deploy(
     cpu_size: str,
     profile: str,
     network: str,
-    function_invoke_opt: str,
     sf_cli_org: Optional[str],
 ):
     from datacustomcode.deploy import (
         COMPUTE_TYPES,
         CodeExtensionMetadata,
+        USE_IN_FEATURE_MAPPING_FOR_CONNECT_API,
         deploy_full,
+        infer_use_in_feature,
     )
     from datacustomcode.token_provider import (
         CredentialsTokenProvider,
@@ -211,15 +234,21 @@ def deploy(
     )
 
     if package_type == "function":
-        if not function_invoke_opt:
+        # Infer use_in_feature from function signature
+        entrypoint_path = os.path.join(path, "entrypoint.py")
+        use_in_feature = infer_use_in_feature(entrypoint_path)
+        if use_in_feature:
+            logger.info(f"Inferred use_in_feature: {use_in_feature}")
+        else:
             click.secho(
-                "Error: Function invoke options are required for function package type",
+                "Error: Could not infer function invoke options. Please provide --use-in-feature",
                 fg="red",
             )
             raise click.Abort()
-        else:
-            function_invoke_options = function_invoke_opt.split(",")
-            metadata.functionInvokeOptions = function_invoke_options
+
+        # Map user-provided feature names to API names
+        mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(use_in_feature, use_in_feature)
+        metadata.functionInvokeOptions = [mapped_feature]
 
     try:
         if sf_cli_org:
@@ -238,19 +267,29 @@ def deploy(
 @click.option(
     "--code-type", default="script", type=click.Choice(["script", "function"])
 )
-def init(directory: str, code_type: str):
+@click.option(
+    "--use-in-feature",
+    default="SearchIndexChunking",
+    help="Feature to invoke the function (only applicable for functions). If not provided, will be inferred from function signature.",
+)
+def init(directory: str, code_type: str, use_in_feature: Optional[str]):
     from datacustomcode.scan import (
         dc_config_json_from_file,
         update_config,
         write_sdk_config,
     )
-    from datacustomcode.template import copy_function_template, copy_script_template
+    from datacustomcode.template import (
+        copy_function_template,
+        copy_script_template,
+    )
 
     click.echo("Copying template to " + click.style(directory, fg="blue", bold=True))
     if code_type == "script":
         copy_script_template(directory)
     elif code_type == "function":
-        copy_function_template(directory)
+        # Default to SearchIndexChunking if not provided
+        feature = use_in_feature
+        copy_function_template(directory, feature)
     entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
     config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
 
@@ -265,6 +304,7 @@ def init(directory: str, code_type: str):
     updated_config_json = update_config(entrypoint_path)
     with open(config_location, "w") as f:
         json.dump(updated_config_json, f, indent=2)
+
     click.echo(
         "Start developing by updating the code in "
         + click.style(entrypoint_path, fg="blue", bold=True)
@@ -275,6 +315,23 @@ def init(directory: str, code_type: str):
         + " to automatically update config.json when you make changes to your code"
     )
 
+    # Generate test.json for functions
+    if code_type == "function":
+        test_json_path = _generate_function_test_file(entrypoint_path)
+        if test_json_path:
+            click.echo(
+                "Generated test file at "
+                + click.style(test_json_path, fg="blue", bold=True)
+            )
+            click.echo(
+                "Test your function locally with "
+                + click.style(
+                    f"datacustomcode run {entrypoint_path} --test_with {test_json_path}",
+                    fg="blue",
+                    bold=True,
+                )
+            )
+
 
 @cli.command()
 @click.argument("filename")
@@ -312,6 +369,12 @@ def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
 @click.option("--config-file", default=None)
 @click.option("--dependencies", default=[], multiple=True)
 @click.option("--profile", default="default")
+@click.option(
+    "--test_with",
+    default=None,
+    type=click.Path(exists=True),
+    help="Path to test JSON file for function testing",
+)
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -322,10 +385,11 @@ def run(
     config_file: Union[str, None],
     dependencies: List[str],
     profile: str,
+    test_with: Optional[str],
     sf_cli_org: Optional[str],
 ):
     from datacustomcode.run import run_entrypoint
 
     run_entrypoint(
-        entrypoint, config_file, dependencies, profile, sf_cli_org=sf_cli_org
+        entrypoint, config_file, dependencies, profile, test_file=test_with, sf_cli_org=sf_cli_org
     )

src/datacustomcode/deploy.py

@@ -65,6 +65,47 @@ def _sanitize_api_name(name: str) -> str:
     return sanitized
 
 
+# Mapping from user-facing feature names to internal API names
+USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
+    "SearchIndexChunking": "UnstructuredChunking",
+}
+
+# Mapping from Pydantic request/response types to feature names
+REQUEST_TYPE_TO_FEATURE = {
+    "SearchIndexChunkingV1Request": "SearchIndexChunking",
+    "SearchIndexChunkingV1Response": "SearchIndexChunking",
+}
+
+def infer_use_in_feature(entrypoint_path: str) -> Union[str, None]:
+    """Infer the use_in_feature from function signature.
+
+    Checks both the request parameter type and return type annotation.
+    Both must map to the same feature for a valid inference.
+
+    Args:
+        entrypoint_path: Path to the entrypoint.py file
+
+    Returns:
+        The feature name if both request and response match, None otherwise
+    """
+    from datacustomcode.function_utils import inspect_function_types
+
+    request_type_name, response_type_name = inspect_function_types(entrypoint_path)
+
+    if not request_type_name or not response_type_name:
+        return None
+
+    # Look up features for both types
+    request_feature = REQUEST_TYPE_TO_FEATURE.get(request_type_name)
+    response_feature = REQUEST_TYPE_TO_FEATURE.get(response_type_name)
+
+    # Both must be present and must match
+    if request_feature and response_feature and request_feature == response_feature:
+        return request_feature
+
+    return None
+
+
 class CodeExtensionMetadata(BaseModel):
     name: str
     version: str

src/datacustomcode/function/feature_types/chunking.py

@@ -28,7 +28,7 @@
 from pydantic import BaseModel, Field
 
 
-class DocElement(BaseModel):
+class SearchIndexDocElement(BaseModel):
     """Document element to be chunked"""
 
     text: str = Field(..., description="Text content to be chunked")
@@ -37,7 +37,7 @@ class DocElement(BaseModel):
     )
 
 
-class ChunkOutput(BaseModel):
+class SearchIndexChunkOutput(BaseModel):
     """Output chunk from the chunking process"""
 
     chunk_id: str = Field(..., description="UUID for this chunk")
@@ -55,20 +55,17 @@ class ChunkOutput(BaseModel):
     )
 
 
-class StatusResponse(BaseModel):
+class SearchIndexStatusResponse(BaseModel):
     """Status response for operation"""
 
     status_type: str = Field(..., description="'success' or 'error'")
     status_message: str = Field(..., description="Human-readable status")
 
 
-class UdsChunkingV1BatchRequest(BaseModel):
+class SearchIndexChunkingV1Request(BaseModel):
     """Batch request for UDS chunking"""
 
-    version: Literal["v1"] = Field(
-        default="v1", description="API version, must be 'v1'"
-    )
-    input: List[DocElement] = Field(
+    input: List[SearchIndexDocElement] = Field(
         ..., min_length=1, description="List of documents (min 1)"
     )
     max_characters: int = Field(..., description="Max chars per chunk (default: 100)")
@@ -77,13 +74,9 @@ class UdsChunkingV1BatchRequest(BaseModel):
     )
 
 
-class UdsChunkingV1BatchResponse(BaseModel):
+class SearchIndexChunkingV1Response(BaseModel):
     """Batch response for UDS chunking"""
-
-    version: Literal["v1"] = Field(
-        default="v1", description="API version, must be 'v1'"
-    )
-    output: List[ChunkOutput] = Field(
+    output: List[SearchIndexChunkOutput] = Field(
         default_factory=list, description="Flat list of chunks from all docs"
     )
-    status: StatusResponse = Field(..., description="Overall operation status")
+    status: SearchIndexStatusResponse = Field(..., description="Overall operation status")