truefoundry 0.6.3rc2__tar.gz → 0.6.4__tar.gz

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.6.3rc2
+Version: 0.6.4
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -30,7 +30,7 @@ Requires-Dist: requirements-parser<0.12.0,>=0.11.0
 Requires-Dist: rich-click<2.0.0,>=1.2.1
 Requires-Dist: rich<14.0.0,>=13.7.1
 Requires-Dist: tqdm<5.0.0,>=4.0.0
-Requires-Dist: truefoundry-sdk==0.0.8
+Requires-Dist: truefoundry-sdk==0.0.10
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truefoundry"
-version = "0.6.3rc2"  # This is a placeholder version, it will be replaced when releasing
+version = "0.6.4"  # This is a placeholder version, it will be replaced when releasing
 requires-python = ">=3.8.1,<3.14"
 description = "TrueFoundry CLI"
 authors = [{ name = "TrueFoundry Team", email = "abhishek@truefoundry.com" }]
@@ -31,7 +31,7 @@ dependencies = [
   "rich>=13.7.1,<14.0.0",
   "rich-click>=1.2.1,<2.0.0",
   "tqdm>=4.0.0,<5.0.0",
-  "truefoundry-sdk==0.0.8",
+  "truefoundry-sdk==0.0.10",
   "typing-extensions>=4.0",
   "urllib3>=1.26.18,<3",
   "yq>=3.1.0,<4.0.0",

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/__init__.py

@@ -14,6 +14,7 @@ from truefoundry.common.warnings import (
     surface_truefoundry_deprecation_warnings,
 )
 from truefoundry.deploy.core import login, logout
+from truefoundry.ml.prompt_utils import render_prompt
 
 surface_truefoundry_deprecation_warnings()
 __all__ = [
@@ -25,6 +26,7 @@ __all__ = [
     "ModelConfiguration",
     "Parameters",
     "PromptVersion",
+    "render_prompt",
     "suppress_truefoundry_deprecation_warnings",
     "SystemMessage",
     "UserMessage",

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/common/utils.py

@@ -93,7 +93,10 @@ def timed_lru_cache(
 
 
 def poll_for_function(
-    func: Callable[..., T], poll_after_secs: int = 5, *args, **kwargs
+    func: Callable[..., T],
+    poll_after_secs: int = 5,
+    *args,
+    **kwargs,
 ) -> Generator[T, None, None]:
     while True:
         yield func(*args, **kwargs)

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/__init__.py

@@ -14,6 +14,7 @@ from truefoundry.deploy._autogen.models import (
     SparkExecutorDynamicScaling,
     SparkExecutorFixedInstances,
     WorkbenchImage,
+    WorkflowAlert,
 )
 from truefoundry.deploy.lib.dao.application import (
     delete_application,
@@ -72,6 +73,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     CUDAVersion,
     DockerFileBuild,
     DynamicVolumeConfig,
+    Email,
     Endpoint,
     GcpTPU,
     GitHelmRepo,
@@ -109,6 +111,10 @@ from truefoundry.deploy.v2.lib.patched_models import (
     Schedule,
     SecretMount,
     ServiceAutoscaling,
+    SlackWebhook,
+    SparkJobJavaEntrypoint,
+    SparkJobPythonEntrypoint,
+    SparkJobScalaEntrypoint,
     SQSInputConfig,
     SQSOutputConfig,
     SQSQueueMetricConfig,

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-03-20T05:09:43+00:00
+#   timestamp: 2025-04-07T11:22:45+00:00
 
 from __future__ import annotations
 
@@ -212,6 +212,17 @@ class DynamicVolumeConfig(BaseModel):
     size: conint(ge=1, le=64000) = Field(..., description="Size of volume in Gi")
 
 
+class Email(BaseModel):
+    type: Literal["email"] = Field(..., description="")
+    notification_channel: constr(min_length=1) = Field(
+        ..., description="Specify the notification channel to send alerts to"
+    )
+    to_emails: List[constr(min_length=1)] = Field(
+        ...,
+        description="List of recipients' email addresses if the notification channel is Email.",
+    )
+
+
 class Endpoint(BaseModel):
     host: constr(
         regex=r"^((([a-zA-Z0-9\-]{1,63}\.)([a-zA-Z0-9\-]{1,63}\.)*([A-Za-z]{1,63}))|(((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)))$"
@@ -271,9 +282,9 @@ class GitSource(BaseModel):
     """
 
     type: Literal["git"] = Field(..., description="")
-    repo_url: constr(
-        regex=r"^(http(s?)://)(github\.com|(.+@)*bitbucket\.org|gitlab\.com|(.*)@dev.azure.com).*$"
-    ) = Field(..., description="The repository URL.")
+    repo_url: constr(regex=r"^(https?://)\S+$") = Field(
+        ..., description="The repository URL."
+    )
     ref: str = Field(..., description="The commit SHA.")
     branch_name: Optional[str] = Field(
         None,
@@ -349,25 +360,6 @@ class Image(BaseModel):
     )
 
 
-class JobAlert(BaseModel):
-    """
-    Describes the configuration for the job alerts
-    """
-
-    notification_channel: constr(min_length=1) = Field(
-        ..., description="Specify the notification channel to send alerts to"
-    )
-    to_emails: Optional[List[constr(min_length=1)]] = Field(
-        None,
-        description="List of recipients' email addresses if the notification channel is Email.",
-    )
-    on_start: bool = Field(False, description="Send an alert when the job starts")
-    on_completion: bool = Field(
-        False, description="Send an alert when the job completes"
-    )
-    on_failure: bool = Field(True, description="Send an alert when the job fails")
-
-
 class Claim(BaseModel):
     key: str
     values: List[str]
@@ -837,6 +829,13 @@ class ServiceAutoscaling(BaseAutoscaling):
     )
 
 
+class SlackWebhook(BaseModel):
+    type: Literal["slack-webhook"] = Field(..., description="")
+    notification_channel: constr(min_length=1) = Field(
+        ..., description="Specify the notification channel to send alerts to"
+    )
+
+
 class SparkDriverConfig(BaseModel):
     ui_endpoint: Endpoint
     resources: Optional[Resources] = None
@@ -857,6 +856,42 @@ class SparkExecutorFixedInstances(BaseModel):
     count: conint(ge=0, le=500) = Field(1, description="Number of instances to start")
 
 
+class SparkJobJavaEntrypoint(BaseModel):
+    type: Literal["java"] = Field(..., description="")
+    main_application_file: str = Field(
+        ..., description="The main application file to be executed by the spark job."
+    )
+    main_class: str = Field(
+        ..., description="The main class to be executed by the spark job."
+    )
+    arguments: Optional[str] = Field(
+        None, description="Arguments to be passed to the main application file."
+    )
+
+
+class SparkJobPythonEntrypoint(BaseModel):
+    type: Literal["python"] = Field(..., description="")
+    main_application_file: str = Field(
+        ..., description="The main application file to be executed by the spark job."
+    )
+    arguments: Optional[str] = Field(
+        None, description="Arguments to be passed to the main application file."
+    )
+
+
+class SparkJobScalaEntrypoint(BaseModel):
+    type: Literal["scala"] = Field(..., description="")
+    main_application_file: str = Field(
+        ..., description="The main application file to be executed by the spark job."
+    )
+    main_class: str = Field(
+        ..., description="The main class to be executed by the spark job."
+    )
+    arguments: Optional[str] = Field(
+        None, description="Arguments to be passed to the main application file."
+    )
+
+
 class StaticVolumeConfig(BaseModel):
     type: Literal["static"] = Field(..., description="Volume Type for the volume.")
     persistent_volume_name: str = Field(
@@ -1174,59 +1209,6 @@ class Helm(BaseModel):
     )
 
 
-class Job(BaseModel):
-    """
-    Describes the configuration for the job
-    """
-
-    type: Literal["job"] = Field(..., description="")
-    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
-        ..., description="Name of the job"
-    )
-    image: Union[Build, Image] = Field(
-        ...,
-        description="Specify whether you want to deploy a Docker image or build and deploy from source code",
-    )
-    trigger: Union[Manual, Schedule] = Field(
-        {"type": "manual"}, description="Specify the trigger"
-    )
-    trigger_on_deploy: bool = Field(
-        False, description="Trigger the job after deploy immediately"
-    )
-    params: Optional[List[Param]] = Field(
-        None, description="Configure params and pass it to create different job runs"
-    )
-    env: Optional[Dict[str, str]] = Field(
-        None,
-        description="Configure environment variables to be injected in the service either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
-    )
-    resources: Optional[Resources] = None
-    alerts: Optional[List[JobAlert]] = Field(
-        None,
-        description="Configure alerts to be sent when the job starts/fails/completes",
-    )
-    retries: conint(ge=0, le=10) = Field(
-        0,
-        description="Specify the maximum number of attempts to retry a job before it is marked as failed.",
-    )
-    timeout: Optional[conint(le=432000, gt=0)] = Field(
-        None, description="Job timeout in seconds."
-    )
-    concurrency_limit: Optional[PositiveInt] = Field(
-        None, description="Number of runs that can run concurrently"
-    )
-    service_account: Optional[str] = Field(None, description="")
-    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
-        None,
-        description="Configure data to be mounted to job pod(s) as a string, secret or volume. [Docs](https://docs.truefoundry.com/docs/mounting-volumes-job)",
-    )
-    labels: Optional[Dict[str, str]] = Field(None, description="")
-    kustomize: Optional[Kustomize] = None
-    workspace_fqn: Optional[str] = Field(
-        None, description="Fully qualified name of the workspace"
-    )
-
-
 class KafkaInputConfig(BaseModel):
     """
     Describes the configuration for the input Kafka worker
@@ -1403,24 +1385,21 @@ class SparkJob(BaseModel):
     name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ..., description="Name of the job"
     )
+    entrypoint: Union[
+        SparkJobPythonEntrypoint, SparkJobScalaEntrypoint, SparkJobJavaEntrypoint
+    ] = Field(..., description="")
     image: Image
     spark_version: str = Field(
         "3.5.2",
         description="Spark version should match the spark version installed in the image.",
     )
-    main_application_file: str = Field(
-        ..., description="The main application file to be executed by the spark job."
-    )
-    arguments: Optional[str] = Field(
-        None, description="Arguments to be passed to the main application file."
-    )
     driver_config: SparkDriverConfig
     executor_config: SparkExecutorConfig
     env: Optional[Dict[str, Any]] = Field(
         None,
         description="Configure environment variables to be injected in the service either as plain text. [Docs](https://docs.truefoundry.com/docs/env-variables)",
     )
-    conf: Optional[Dict[str, Any]] = Field(
+    spark_conf: Optional[Dict[str, Any]] = Field(
         None,
         description="Extra configuration properties to be passed to the spark job. [Docs](https://spark.apache.org/docs/latest/configuration.html)",
     )
@@ -1468,6 +1447,18 @@ class WorkerConfig(BaseModel):
     )
 
 
+class WorkflowAlert(BaseModel):
+    """
+    Describes the configuration for the workflow alerts
+    """
+
+    notification_target: Optional[Union[Email, SlackWebhook]] = None
+    on_completion: bool = Field(
+        False, description="Send an alert when the job completes"
+    )
+    on_failure: bool = Field(True, description="Send an alert when the job fails")
+
+
 class BaseService(BaseModel):
     name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
@@ -1516,6 +1507,26 @@ class FlyteTaskTemplate(BaseModel):
     custom: FlyteTaskCustom
 
 
+class JobAlert(BaseModel):
+    """
+    Describes the configuration for the job alerts
+    """
+
+    notification_channel: Optional[constr(min_length=1)] = Field(
+        None, description="Specify the notification channel to send alerts to"
+    )
+    to_emails: Optional[List[constr(min_length=1)]] = Field(
+        None,
+        description="List of recipients' email addresses if the notification channel is Email.",
+    )
+    notification_target: Optional[Union[Email, SlackWebhook]] = None
+    on_start: bool = Field(False, description="Send an alert when the job starts")
+    on_completion: bool = Field(
+        False, description="Send an alert when the job completes"
+    )
+    on_failure: bool = Field(True, description="Send an alert when the job fails")
+
+
 class Service(BaseService):
     """
     Describes the configuration for the service
@@ -1557,6 +1568,59 @@ class FlyteTask(BaseModel):
     description: Optional[Any] = None
 
 
+class Job(BaseModel):
+    """
+    Describes the configuration for the job
+    """
+
+    type: Literal["job"] = Field(..., description="")
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
+        ..., description="Name of the job"
+    )
+    image: Union[Build, Image] = Field(
+        ...,
+        description="Specify whether you want to deploy a Docker image or build and deploy from source code",
+    )
+    trigger: Union[Manual, Schedule] = Field(
+        {"type": "manual"}, description="Specify the trigger"
+    )
+    trigger_on_deploy: bool = Field(
+        False, description="Trigger the job after deploy immediately"
+    )
+    params: Optional[List[Param]] = Field(
+        None, description="Configure params and pass it to create different job runs"
+    )
+    env: Optional[Dict[str, str]] = Field(
+        None,
+        description="Configure environment variables to be injected in the service either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
+    )
+    resources: Optional[Resources] = None
+    alerts: Optional[List[JobAlert]] = Field(
+        None,
+        description="Configure alerts to be sent when the job starts/fails/completes",
+    )
+    retries: conint(ge=0, le=10) = Field(
+        0,
+        description="Specify the maximum number of attempts to retry a job before it is marked as failed.",
+    )
+    timeout: Optional[conint(le=432000, gt=0)] = Field(
+        None, description="Job timeout in seconds."
+    )
+    concurrency_limit: Optional[PositiveInt] = Field(
+        None, description="Number of runs that can run concurrently"
+    )
+    service_account: Optional[str] = Field(None, description="")
+    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
+        None,
+        description="Configure data to be mounted to job pod(s) as a string, secret or volume. [Docs](https://docs.truefoundry.com/docs/mounting-volumes-job)",
+    )
+    labels: Optional[Dict[str, str]] = Field(None, description="")
+    kustomize: Optional[Kustomize] = None
+    workspace_fqn: Optional[str] = Field(
+        None, description="Fully qualified name of the workspace"
+    )
+
+
 class Workflow(BaseModel):
     """
     Describes the configuration for the worflow
@@ -1575,6 +1639,7 @@ class Workflow(BaseModel):
     flyte_entities: Optional[List[Union[FlyteTask, FlyteWorkflow, FlyteLaunchPlan]]] = (
         Field(None, description="")
     )
+    alerts: Optional[List[WorkflowAlert]] = Field(None, description="")
 
 
 class ApplicationSet(BaseModel):

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -14,7 +14,10 @@ from rich.status import Status
 from tqdm import tqdm
 from tqdm.utils import CallbackIOWrapper
 
-from truefoundry.common.constants import VERSION_PREFIX
+from truefoundry.common.constants import (
+    SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+    VERSION_PREFIX,
+)
 from truefoundry.common.request_utils import request_handling
 from truefoundry.common.servicefoundry_client import (
     ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
@@ -247,7 +250,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
-        sio = socketio.Client(request_timeout=60)
+        sio = socketio.Client(request_timeout=60, reconnection_attempts=10)
         callback.print_line("Waiting for the task to start...")
         next_log_start_timestamp = query_dict.get("startTs")
 
@@ -286,6 +289,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 transports="websocket",
                 headers=self._get_headers(),
                 socketio_path=socketio_path,
+                retry=True,
             )
             return sio
 
@@ -298,10 +302,15 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
 
     @check_min_cli_version
     def get_deployment_statuses(
-        self, application_id: str, deployment_id: str
+        self,
+        application_id: str,
+        deployment_id: str,
+        retry_count: int = SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
     ) -> List[AppDeploymentStatusResponse]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}/statuses"
-        response = session_with_retries().get(url, headers=self._get_headers())
+        response = session_with_retries(retries=retry_count).get(
+            url, headers=self._get_headers()
+        )
         response_data = request_handling(response)
         return parse_obj_as(List[AppDeploymentStatusResponse], response_data)
 

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/v2/lib/deploy.py

@@ -107,12 +107,20 @@ def _tail_build_logs(build_response: BuildResponse) -> socketio.Client:
 
 def _deploy_wait_handler(  # noqa: C901
     deployment: Deployment,
+    tail_logs: bool = True,
 ) -> Optional[DeploymentTransitionStatus]:
+    tail_logs_or_polling_status_message = (
+        "You can press Ctrl + C to exit the tailing of build logs "
+    )
+    if not tail_logs:
+        tail_logs_or_polling_status_message = (
+            "you can press Ctrl + C to exit the polling of deployment status "
+        )
     _log_application_dashboard_url(
         deployment=deployment,
         log_message=(
             "You can track the progress below or on the dashboard:- '%s'\n"
-            "You can press Ctrl + C to exit the tailing of build logs "
+            f"{tail_logs_or_polling_status_message}"
             "and deployment will continue on the server"
         ),
     )
@@ -130,6 +138,7 @@ def _deploy_wait_handler(  # noqa: C901
             poll_after_secs=poll_interval_seconds,
             application_id=deployment.applicationId,
             deployment_id=deployment.id,
+            retry_count=10,
         ):
             if len(deployment_statuses) == 0:
                 logger.warning("Did not receive any deployment status")
@@ -157,12 +166,15 @@ def _deploy_wait_handler(  # noqa: C901
                 latest_deployment_status.transition
                 == DeploymentTransitionStatus.BUILDING
             ):
-                if not socket:
-                    build_responses = client.get_deployment_build_response(
-                        application_id=deployment.applicationId,
-                        deployment_id=deployment.id,
-                    )
-                    socket = _tail_build_logs(build_responses[0])
+                if tail_logs and not socket:
+                    try:
+                        build_responses = client.get_deployment_build_response(
+                            application_id=deployment.applicationId,
+                            deployment_id=deployment.id,
+                        )
+                        socket = _tail_build_logs(build_responses[0])
+                    except Exception as e:
+                        logger.error("Error tailing build logs: %s", e)
 
             time_elapsed = time.monotonic() - start_time
             if time_elapsed > total_timeout_time:

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/v2/lib/deploy_workflow.py

@@ -25,15 +25,17 @@ from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
 from truefoundry.deploy.lib.dao.workspace import get_workspace_by_fqn
-from truefoundry.deploy.lib.model.entity import Deployment
+from truefoundry.deploy.lib.model.entity import Deployment, DeploymentTransitionStatus
+from truefoundry.deploy.v2.lib.deploy import _deploy_wait_handler
 from truefoundry.deploy.v2.lib.source import (
     local_source_to_remote_source,
 )
 from truefoundry.logger import logger
 from truefoundry.pydantic_v1 import ValidationError
 from truefoundry.workflow.workflow import (
+    TRUEFOUNDRY_ALERTS_CONFIG,
     TRUEFOUNDRY_LAUNCH_PLAN_NAME,
-    execution_config_store,
+    truefoundry_config_store,
 )
 
 
@@ -173,7 +175,7 @@ def _validate_workflow_entities(  # noqa: C901
     # validate that all inputs have default values for cron workflows
     for launch_plan in launch_plans:
         if (
-            execution_config_store.get(launch_plan.spec.workflow_id.name)
+            truefoundry_config_store.get(launch_plan.spec.workflow_id.name)
             and launch_plan.id.name == TRUEFOUNDRY_LAUNCH_PLAN_NAME
         ):
             workflow_inputs = launch_plan.spec.default_inputs.parameters
@@ -236,7 +238,7 @@ def _generate_manifest_for_workflow(
 
             # this is the case when someone has a cron schedule. and this line is for handling default launch plan in this case.
             if (
-                execution_config_store.get(workflow_name)
+                truefoundry_config_store.get(workflow_name)
                 and workflow_name == entity.id.name
             ):
                 continue
@@ -262,6 +264,14 @@ def _generate_manifest_for_workflow(
 
         workflow.flyte_entities.append(message_dict)
 
+    alerts_config = truefoundry_config_store.get(TRUEFOUNDRY_ALERTS_CONFIG)
+    if alerts_config:
+        if not workflow.alerts:
+            workflow.alerts = alerts_config
+        else:
+            logger.warning(
+                "Alerts are configured in both workflow decorator as well as in deployment config. Alerts configured in workflow decorator will be ignored."
+            )
     # this step is just to verify if pydantic model is still valid after adding flyte_entities
     autogen_models.Workflow.validate({**workflow.dict()})
 
@@ -288,7 +298,7 @@ def deploy_workflow(
 
     # we need to rest the execution config store as it is a global variable and we don't want to keep the cron execution config for next workflow
     # this is only needed for notebook environment
-    execution_config_store.reset()
+    truefoundry_config_store.reset()
 
     workspace_id = get_workspace_by_fqn(workspace_fqn).id
 
@@ -314,4 +324,17 @@ def deploy_workflow(
         deployment.fqn,
     )
     deployment_url = f"{client.tfy_host.strip('/')}/applications/{deployment.applicationId}?tab=deployments"
+    if wait:
+        try:
+            last_status_printed = _deploy_wait_handler(
+                deployment=deployment, tail_logs=False
+            )
+            if not last_status_printed or DeploymentTransitionStatus.is_failure_state(
+                last_status_printed
+            ):
+                deployment_tab_url = f"{client.tfy_host.strip('/')}/applications/{deployment.applicationId}?tab=deployments"
+                message = f"Deployment Failed. Please refer to the logs for additional details - {deployment_tab_url}"
+                sys.exit(message)
+        except KeyboardInterrupt:
+            logger.info("Ctrl-c executed. The deployment will still continue.")
     logger.info("You can find the application on the dashboard:- '%s'", deployment_url)

{truefoundry-0.6.3rc2 → truefoundry-0.6.4}/truefoundry/deploy/v2/lib/patched_models.py

@@ -563,3 +563,23 @@ class TaskPythonBuild(models.TaskPythonBuild, PatchedModelBase):
 
 # Deprecated aliases, kept for backward compatibility
 TruefoundryArtifactSource = TrueFoundryArtifactSource
+
+
+class Email(models.Email, PatchedModelBase):
+    type: Literal["email"] = "email"
+
+
+class SlackWebhook(models.SlackWebhook, PatchedModelBase):
+    type: Literal["slack-webhook"] = "slack-webhook"
+
+
+class SparkJobScalaEntrypoint(models.SparkJobScalaEntrypoint, PatchedModelBase):
+    type: Literal["scala"] = "scala"
+
+
+class SparkJobPythonEntrypoint(models.SparkJobPythonEntrypoint, PatchedModelBase):
+    type: Literal["python"] = "python"
+
+
+class SparkJobJavaEntrypoint(models.SparkJobJavaEntrypoint, PatchedModelBase):
+    type: Literal["java"] = "java"

truefoundry-0.6.4/truefoundry/ml/prompt_utils.py

@@ -0,0 +1,65 @@
+import re
+from typing import Any, Dict, Optional
+
+from truefoundry_sdk import ChatPromptManifest
+
+_VAR_PATTERN = re.compile(r"\{\{(\w+)\}\}")
+
+
+def render_text(template: str, variables: Dict[str, str]) -> str:
+    return _VAR_PATTERN.sub(
+        lambda m: (variables.get(m.group(1), f"{{{{{m.group(1)}}}}}")), template
+    )
+
+
+def render_prompt(
+    prompt_template: ChatPromptManifest, variables: Optional[Dict[str, str]] = None
+) -> Dict[str, Any]:
+    """
+    Renders a prompt template with the provided variables.
+    Args:
+        prompt_template (ChatPromptManifest): The prompt template to render.
+        variables (Optional[Dict[str, str]]): A dictionary of variables to replace in the template.
+            If None, the default variables from the prompt template will be used.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the rendered messages and model configuration.
+
+    """
+    # Render the messages in the prompt template with the provided variables and default variables
+    rendered_messages = []
+    variables = {**(prompt_template.variables or {}), **(variables or {})}
+    messages = [message.copy(deep=True) for message in prompt_template.messages]
+
+    # Render the messages with the provided variables, if any
+    if variables:
+        for msg in messages:
+            if isinstance(msg.content, str):
+                msg.content = render_text(template=msg.content, variables=variables)
+            elif isinstance(msg.content, list):
+                for part in msg.content:
+                    if part.type == "image_url" and isinstance(part.image_url.url, str):
+                        part.image_url.url = render_text(
+                            part.image_url.url, variables=variables
+                        )
+                    elif part.type == "text" and isinstance(part.text, str):
+                        part.text = render_text(part.text, variables=variables)
+
+    rendered_messages = [message.dict() for message in messages]
+
+    # Merge parameters from model_configuration and extra_parameters
+    model_configuration = prompt_template.model_configuration
+    parameters = {
+        **(
+            model_configuration.parameters.dict()
+            if model_configuration.parameters
+            else {}
+        ),
+        **(model_configuration.extra_parameters or {}),
+    }
+
+    return {
+        "messages": rendered_messages,
+        "model": model_configuration.model,
+        "parameters": parameters,
+    }