truefoundry 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

truefoundry/ml/autogen/entities/artifacts.py

@@ -0,0 +1,286 @@
+# generated by datamodel-codegen:
+#   filename:  artifacts.json
+#   timestamp: 2024-09-12T13:40:36+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from truefoundry.pydantic_v1 import BaseModel, Field, constr
+
+
+class Agent(BaseModel):
+    type: Literal["agent"] = Field(..., description="+value=agent")
+    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
+        ..., description="Names cannot repeat accross agent\n+uiType=Hidden"
+    )
+    available_tools: List[constr(min_length=1)] = Field(
+        ..., description="+sort=30\n+uiType=Hidden"
+    )
+    goal: constr(min_length=1, max_length=128) = Field(
+        ..., description="+sort=10\n+uiType=TextArea"
+    )
+    instruction: constr(min_length=1, max_length=2620) = Field(
+        ...,
+        description='`instruction` is the system prompt for now. (2.5 * 1024)\n+sort=20\n+uiType=AgentInstructions\n+uiProps={"helpText":"Use the syntax ${Tool FQN} to reference a tool, and ${AGENT FQN} to reference another agent"}',
+    )
+    model_id: constr(min_length=1) = Field(
+        ...,
+        description='+sort=40\n+uiType=EnabledModelSelector\n+uiProps={"searchable":true,"modelType":"chat","providerType":"openai"}',
+    )
+
+
+class Method(str, Enum):
+    """
+    +sort=50
+    +uiType=Hidden
+    """
+
+    get = "get"
+    post = "post"
+    put = "put"
+    delete = "delete"
+    patch = "patch"
+
+
+class AgentWithFQN(Agent):
+    id: str
+    fqn: str
+
+
+class MimeType(str, Enum):
+    """
+    +label=MIME Type
+    +usage=MIME type of the content
+    """
+
+    text_plain = "text/plain"
+    application_json = "application/json"
+    image_png = "image/png"
+    image_jpeg = "image/jpeg"
+    application_x_directory = "application/x-directory"
+
+
+class BlobStorageReference(BaseModel):
+    """
+    +usage=Blob Storage Location
+    +label=Blob Storage Location
+    +docs=Defines the structure for blob storage content references, including type, path, and mime_type
+    """
+
+    type: Literal["blob-storage"]
+    path: constr(regex=r"^.{1,}$") = Field(
+        ..., description="+label=Path\n+usage=Path to the content in blob storage"
+    )
+    mime_type: MimeType = Field(
+        ..., description="+label=MIME Type\n+usage=MIME type of the content"
+    )
+
+
+class ImageUrl(BaseModel):
+    """
+    +label=URL for the image
+    +usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL
+    """
+
+    url: Union[
+        constr(
+            regex=r"^\b((https?://)?(?:www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}(?:[-a-zA-Z0-9()@:%_\+.~#?&/=]*)|data:image/[a-zA-Z]+;base64,[a-zA-Z0-9+/=\s]+)$"
+        ),
+        BlobStorageReference,
+    ] = Field(
+        ...,
+        description="+label=URL for the image\n+usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL",
+    )
+
+
+class ImageContentPart(BaseModel):
+    """
+    +usage=Image URL
+    +label=Image URL
+    """
+
+    type: Literal["image_url"]
+    image_url: ImageUrl = Field(
+        ...,
+        description="+label=URL for the image\n+usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL",
+    )
+
+
+class Parameters(BaseModel):
+    """
+    +usage=Parameters for the provider
+    +label=Parameters
+    +docs=Key-value pairs to store additional parameters for the provider
+    """
+
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_k: Optional[float] = None
+    top_p: Optional[float] = None
+    stop: Optional[Union[List[str], str]] = None
+
+
+class ModelConfiguration(BaseModel):
+    """
+    +label=Model Configuration
+    +icon=fa-cogs:#326ce5
+    +message=Configuration details for the provider and corresponding config
+    +usage=Define the provider and its configuration
+    +docs=Configuration settings specific to the provider, including model settings and other parameters.
+    """
+
+    provider: str = Field(
+        ...,
+        description='+label=Provider Name\n+usage=Name of the provider, must be non-empty, e.g., "openai", "google_gemini"',
+    )
+    model: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
+        ..., description="+label=Model name\n+usage=Model name, must be non-empty"
+    )
+    parameters: Optional[Parameters] = Field(
+        None,
+        description="+usage=Parameters for the provider\n+label=Parameters\n+docs=Key-value pairs to store additional parameters for the provider",
+    )
+    extra_parameters: Optional[Dict[str, Any]] = Field(
+        None,
+        description="+usage=Extra parameters for the provider\n+label=Extra Parameters\n+docs=Additional parameters for the provider",
+    )
+
+
+class SystemMessage(BaseModel):
+    """
+    +usage=System message
+    +docs=Defines the structure of a system message, including role and content
+    +label=System Message
+    """
+
+    role: Literal["system"]
+    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
+        ..., description="+label=Content\n+usage=Text content for the system message"
+    )
+    name: Optional[str] = Field(
+        None, description="+label=Name\n+usage=Name of the system"
+    )
+
+
+class TextContentPart(BaseModel):
+    """
+    +usage=Text content
+    +label=Text content
+    """
+
+    type: Literal["text"]
+    text: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference]
+
+
+class UserMessage(BaseModel):
+    """
+    +usage=User message
+    +docs=Defines the structure of a user message, including role and content
+    +label=User Message
+    """
+
+    role: Literal["user"]
+    content: Union[
+        constr(regex=r"^.[\s\S]*$"),
+        BlobStorageReference,
+        List[Union[TextContentPart, ImageContentPart]],
+    ] = Field(
+        ..., description="+label=Content\n+usage=Text content for the user message"
+    )
+    name: Optional[str] = Field(
+        None, description="+label=Name\n+usage=Name of the user message"
+    )
+
+
+class AgentOpenAPITool(BaseModel):
+    type: Literal["openapi-tool"] = Field(..., description="+value=openapi-tool")
+    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
+        ..., description="Names cannot repeat accross tools\n+uiType=Hidden"
+    )
+    description: constr(min_length=1, max_length=128) = Field(
+        ..., description="+sort=10\n+uiType=TextArea"
+    )
+    openapi_spec: Union[BlobStorageReference, Dict[str, Any]] = Field(
+        ..., description="+sort=20\n+uiType=OpenapiSchema"
+    )
+    base_url: constr(
+        regex=r"^(https?://)?(?:www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}(?:[-a-zA-Z0-9()@:%_\+.~#?&/=]*)$"
+    ) = Field(..., description="+sort=30")
+    path: constr(min_length=1) = Field(
+        ...,
+        description="+sort=40\n+uiType=MethodPathSelector\n+label=Method + Path\n+usage=Select one Method and Path for this tool",
+    )
+    method: Method = Field(..., description="+sort=50\n+uiType=Hidden")
+    headers: Optional[Dict[str, str]] = Field(
+        None,
+        description='+sort=60\n+uiType=KV\n+uiProps={"allowSecrets":true,"secretConfig":{"enableNew":true,"hideOptions":true}}',
+    )
+
+
+class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
+    id: str
+    fqn: str
+
+
+class AssistantMessage(BaseModel):
+    """
+    +usage=Assistant message
+    +docs=Defines the structure of an assistant message, including role and content
+    +label=Assistant Message
+    """
+
+    role: Literal["assistant"]
+    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
+        ..., description="+label=Content\n+usage=Text content for the assistant message"
+    )
+    name: Optional[str] = Field(
+        None, description="+label=Name\n+usage=Name of the assistant message"
+    )
+
+
+class BasePrompt(BaseModel):
+    """
+    Main prompt structure
+    """
+
+    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9-_]{0,254}[a-zA-Z0-9]$") = Field(
+        ...,
+        description="+label=Prompt Name\n+icon=fa-desktop:#326ce5\n+message=Alphanumeric word, may contain '-' with a maximum length of 256 characters\n+sort=1000\n+usage=Name of the Prompt. This uniquely identifies it in the workspace.\n+docs=The unique name for the prompt, consisting of alphanumeric characters and dashes, max length 100.",
+    )
+    description: Optional[constr(min_length=1, max_length=512)] = Field(
+        None,
+        description="+label=Description\n+icon=fa-desktop:#326ce5\n+message=Description in a maximum of 512 characters\n+sort=2000\n+usage=Description of the prompt\n+docs=An optional description for the prompt, with a maximum length of 512 characters.\n+optional",
+    )
+    model_configuration: ModelConfiguration
+    metadata: Dict[str, str] = Field(
+        ...,
+        description="+label=Metadata\n+icon=fa-info-circle:#326ce5\n+message=Additional metadata for the prompt\n+sort=7000\n+usage=Define additional metadata for the prompt\n+docs=Key-value pairs to store additional metadata related to the prompt.",
+    )
+
+
+class ChatPrompt(BasePrompt):
+    type: Literal["chat_prompt"] = Field(
+        ...,
+        description='+label=Type\n+icon=fa-desktop:#326ce5\n+message=Type of the prompt\n+sort=3000\n+usage=Type of the prompt\n+docs=Type of the prompt, set to "chat_prompt"\n+default="chat_prompt"',
+    )
+    messages: List[Union[SystemMessage, AssistantMessage, UserMessage]] = Field(
+        ...,
+        description="+sort=4000\n+usage=Chat completion messages\n+label=Messages in the chat conversation\n+message=Chat completion messages\n+usage=List of messages in the chat conversation, must be non-empty\n+docs=Messages that define the chat conversation, including system, assistant, and user messages.",
+    )
+    variables: Optional[Dict[str, Optional[str]]] = Field(
+        None,
+        description="+label=Variables\n+usage=Variables for the chat completion messages to be used in the prompt messages\n+sort=5000",
+    )
+
+
+class AgentApp(BaseModel):
+    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
+    tools: List[AgentOpenAPIToolWithFQN]
+    agents: List[AgentWithFQN]
+    root_agent: constr(min_length=1)
+
+
+class Artifacts(BaseModel):
+    __root__: Union[ChatPrompt, AgentOpenAPITool, Agent, AgentApp]

truefoundry/ml/clients/servicefoundry_client.py

@@ -12,15 +12,13 @@ from truefoundry.ml.clients.entities import (
 )
 from truefoundry.ml.clients.utils import http_request_safe
 from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.ml.run_utils import append_servicefoundry_path_to_tracking_uri
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
+    # TODO (chiragjn): Rename tracking_uri to tfy_host
     def __init__(self, tracking_uri: str, token: Optional[str] = None):
         super().__init__(base_url=tracking_uri)
-        self.host_creds = HostCreds(
-            host=append_servicefoundry_path_to_tracking_uri(tracking_uri), token=token
-        )
+        self.host_creds = HostCreds(host=self._api_server_url, token=token)
 
     def get_integration_from_id(self, integration_id: str):
         integration_id = integration_id or ""

truefoundry/ml/log_types/artifacts/artifact.py

@@ -5,7 +5,7 @@ import os
 import tempfile
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tuple, Union
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
@@ -25,6 +25,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     InternalMetadata as InternalMetadataDto,
 )
+from truefoundry.ml.autogen.entities.artifacts import ChatPrompt
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import INTERNAL_METADATA_PATH
 from truefoundry.ml.log_types.artifacts.utils import (
@@ -351,6 +352,63 @@ class ArtifactVersion:
         self._set_mutable_attrs()
 
 
+class ChatPromptVersion(ArtifactVersion):
+    def __init__(self, artifact_version: ArtifactVersionDto, artifact: ArtifactDto):
+        if artifact.type != ArtifactType.CHAT_PROMPT:
+            raise ValueError(
+                f"{artifact_version.fqn!r} is not a chat prompt type artifact"
+            )
+        super().__init__(
+            artifact_version=artifact_version,
+            artifact=artifact,
+        )
+        self._chat_prompt = ChatPrompt.parse_obj(
+            artifact_version.internal_metadata.to_dict()
+        )
+
+    @classmethod
+    def from_fqn(cls, fqn: str) -> "ChatPromptVersion":
+        api_client = _get_api_client()
+        mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
+        _artifact_version = mlfoundry_artifacts_api.get_artifact_version_by_fqn_get(
+            fqn=fqn
+        )
+        artifact_version = _artifact_version.artifact_version
+        _artifact = mlfoundry_artifacts_api.get_artifact_by_id_get(
+            id=artifact_version.artifact_id
+        )
+        return cls(
+            artifact_version=_artifact_version.artifact_version,
+            artifact=_artifact.artifact,
+        )
+
+    @property
+    def model(self) -> str:
+        return self._chat_prompt.model_configuration.model
+
+    @property
+    def provider(self) -> str:
+        return self._chat_prompt.model_configuration.provider
+
+    @property
+    def messages(self) -> List[Dict[str, Any]]:
+        return [message.dict() for message in self._chat_prompt.messages]
+
+    @property
+    def parameters(self) -> Dict[str, Any]:
+        _parameters = self._chat_prompt.model_configuration.parameters
+        return _parameters.dict(exclude_unset=True) if _parameters else {}
+
+    @property
+    def extra_parameters(self) -> Dict[str, Any]:
+        _extra_parameters = self._chat_prompt.model_configuration.extra_parameters
+        return _extra_parameters.dict(exclude_unset=True) if _extra_parameters else {}
+
+    @property
+    def variables(self) -> Dict[str, Any]:
+        return self._chat_prompt.variables or {}
+
+
 def _log_artifact_version_helper(
     run: "MlFoundryRun",
     name: str,

truefoundry/ml/mlfoundry_api.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 import coolname
 import pandas as pd
 
+from truefoundry.common.utils import relogin_error_message
 from truefoundry.ml import constants, env_vars
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
@@ -34,7 +35,11 @@ from truefoundry.ml.clients.servicefoundry_client import ServiceFoundryServiceCl
 from truefoundry.ml.enums import ModelFramework, ViewType
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
-from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVersion
+from truefoundry.ml.log_types.artifacts.artifact import (
+    ArtifactPath,
+    ArtifactVersion,
+    ChatPromptVersion,
+)
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
 from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
@@ -209,8 +214,9 @@ class MlFoundry:
         session = get_active_session()
         if session is None:
             raise MlFoundryException(
-                "No active session found. Perhaps you are not logged in?\n"
-                "Please log in using `tfy login [--host HOST] --relogin"
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
             )
         servicefoundry_client = ServiceFoundryServiceClient(
             tracking_uri=self.get_tracking_uri(),
@@ -1554,6 +1560,45 @@ class MlFoundry:
             if not datasets or not page_token:
                 done = True
 
+    def get_chat_prompt_version_by_fqn(self, fqn: str) -> ChatPromptVersion:
+        """
+        Get the prompt by prompt version FQN
+        Args:
+            fqn (str): Fully qualified name of the chat prompt version.
+        Returns:
+            ChatPromptVersion: An instance of the ChatPromptVersion class with the fetched prompt details.
+
+        Examples:
+            ```python
+            from truefoundry.ml import get_client
+            client = get_client()
+
+            # Get the chat prompt by FQN
+            fqn = 'chat_prompt:truefoundry/prompt-demo/demo-prompt-1:2'
+            chat_prompt = client.get_chat_prompt_version_by_fqn(fqn=fqn)
+
+            # Check the chat prompt values
+            print(chat_prompt)
+
+            # Use with OpenAI client
+            from openai import OpenAI
+            openai_client = OpenAI(api_key="api_key")
+
+            # Make the OpenAI call
+            response = openai_client.chat.completions.create(
+                model=chat_prompt.model,
+                messages=chat_prompt.messages,
+                **chat_prompt.parameters,
+                **chat_prompt.extra_parameters
+            )
+
+            # Extract the content from the response
+            response_content = response.choices[0].message.content
+            print(response_content)
+            ```
+        """
+        return ChatPromptVersion.from_fqn(fqn=fqn)
+
 
 def get_client() -> MlFoundry:
     """Initializes and returns the mlfoundry client.

truefoundry/ml/mlfoundry_run.py

@@ -19,6 +19,7 @@ from typing import (
 from urllib.parse import urljoin, urlsplit
 
 from truefoundry import version
+from truefoundry.common.utils import relogin_error_message
 from truefoundry.ml import constants, enums
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
@@ -51,6 +52,8 @@ from truefoundry.ml.run_utils import ParamsType, flatten_dict, process_params
 from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client, get_active_session
 from truefoundry.ml.validation_utils import (
     MAX_ENTITY_KEY_LENGTH,
+    MAX_METRICS_PER_BATCH,
+    MAX_PARAMS_TAGS_PER_BATCH,
     _validate_batch_log_data,
 )
 
@@ -222,8 +225,9 @@ class MlFoundryRun:
         session = get_active_session()
         if session is None:
             raise MlFoundryException(
-                "No active session found. Perhaps you are not logged in?\n"
-                "Please log in using `tfy login [--host HOST] --relogin"
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
             )
         base_url = "{uri.scheme}://{uri.netloc}/".format(
             uri=urlsplit(session.tracking_uri)
@@ -636,12 +640,16 @@ class MlFoundryRun:
             return
 
         try:
-            _validate_batch_log_data(metrics=metrics, params=[], tags=[])
-            self._runs_api.log_run_batch_post(
-                log_batch_request_dto=LogBatchRequestDto(
-                    run_id=self.run_id, metrics=metrics, params=[], tags=[]
+            for i in range(0, len(metrics), MAX_METRICS_PER_BATCH):
+                metrics_batch = metrics[i : i + MAX_METRICS_PER_BATCH]
+
+                _validate_batch_log_data(metrics=metrics_batch, params=[], tags=[])
+
+                self._runs_api.log_run_batch_post(
+                    log_batch_request_dto=LogBatchRequestDto(
+                        run_id=self.run_id, metrics=metrics_batch, params=[], tags=[]
+                    )
                 )
-            )
         except Exception as e:
             raise MlFoundryException(str(e)) from e
 
@@ -723,12 +731,17 @@ class MlFoundryRun:
             if len(params) == 0:
                 logger.warning("Cannot log empty params dictionary")
 
-            _validate_batch_log_data(metrics=[], params=params, tags=[])
-            self._runs_api.log_run_batch_post(
-                log_batch_request_dto=LogBatchRequestDto(
-                    run_id=self.run_id, metrics=[], params=params, tags=[]
+            for i in range(0, len(params), MAX_PARAMS_TAGS_PER_BATCH):
+                params_batch = params[i : i + MAX_PARAMS_TAGS_PER_BATCH]
+
+                _validate_batch_log_data(metrics=[], params=params_batch, tags=[])
+                logger.debug("Logging parameters: %s", params_batch)
+
+                self._runs_api.log_run_batch_post(
+                    log_batch_request_dto=LogBatchRequestDto(
+                        run_id=self.run_id, metrics=[], params=params_batch, tags=[]
+                    )
                 )
-            )
         except Exception as e:
             raise MlFoundryException(str(e)) from e
         logger.info("Parameters logged successfully")
@@ -765,11 +778,15 @@ class MlFoundryRun:
             tags_arr = [
                 RunTagDto(key=key, value=str(value)) for key, value in tags.items()
             ]
-            self._runs_api.log_run_batch_post(
-                log_batch_request_dto=LogBatchRequestDto(
-                    run_id=self.run_id, metrics=[], params=[], tags=tags_arr
+            for i in range(0, len(tags_arr), MAX_PARAMS_TAGS_PER_BATCH):
+                tags_batch = tags_arr[i : i + MAX_PARAMS_TAGS_PER_BATCH]
+
+                _validate_batch_log_data(metrics=[], params=[], tags=tags_batch)
+                self._runs_api.log_run_batch_post(
+                    log_batch_request_dto=LogBatchRequestDto(
+                        run_id=self.run_id, metrics=[], params=[], tags=tags_batch
+                    )
                 )
-            )
         except Exception as e:
             raise MlFoundryException(str(e)) from e
         logger.info("Tags set successfully")

truefoundry/ml/run_utils.py

@@ -1,9 +1,7 @@
 import argparse
 import importlib
 import json
-import os
 from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
-from urllib.parse import urljoin, urlsplit
 
 import numpy as np
 
@@ -21,18 +19,6 @@ def get_module(
             raise MlFoundryException(msg) from ex
 
 
-def append_path_to_rest_tracking_uri(tracking_uri: str):
-    if urlsplit(tracking_uri).netloc.startswith("localhost"):
-        return tracking_uri
-    return urljoin(tracking_uri, "/api/ml")
-
-
-def append_servicefoundry_path_to_tracking_uri(tracking_uri: str):
-    if urlsplit(tracking_uri).netloc.startswith("localhost"):
-        return os.getenv("SERVICEFOUNDRY_SERVER_URL")
-    return urljoin(tracking_uri, "/api/svc")
-
-
 class NumpyEncoder(json.JSONEncoder):
     """Special json encoder for numpy types"""
 

truefoundry/ml/session.py

@@ -10,6 +10,7 @@ from truefoundry.common.credential_provider import (
 )
 from truefoundry.common.entities import Token, UserInfo
 from truefoundry.common.request_utils import urllib3_retry
+from truefoundry.common.utils import get_tfy_servers_config, relogin_error_message
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     Configuration,
@@ -17,9 +18,6 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
 from truefoundry.ml.clients.entities import HostCreds
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
-from truefoundry.ml.run_utils import (
-    append_path_to_rest_tracking_uri,
-)
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -83,6 +81,7 @@ class Session:
         self._assert_not_closed()
         return self._user_info
 
+    # TODO (chiragjn): Rename tracking_uri to tfy_host
     @property
     def tracking_uri(self) -> str:
         return self._cred_provider.base_url
@@ -97,7 +96,7 @@ class Session:
         )
 
     def get_host_creds(self) -> HostCreds:
-        tracking_uri = append_path_to_rest_tracking_uri(self._cred_provider.base_url)
+        tracking_uri = get_tfy_servers_config(self.tracking_uri).mlfoundry_server_url
         return HostCreds(
             host=tracking_uri, token=self._cred_provider.token.access_token
         )
@@ -120,8 +119,9 @@ def _get_api_client(
             return ApiClient()
         else:
             raise MlFoundryException(
-                "No active session found. Perhaps you are not logged in?\n"
-                "Please log in using `tfy login [--host HOST] --relogin"
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
             )
 
     creds = session.get_host_creds()
@@ -144,8 +144,9 @@ def init_session() -> Session:
                 break
         if final_cred_provider is None:
             raise MlFoundryException(
-                "Please login using `tfy login` command "
-                "or `truefoundry.login()` function call"
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
             )
         new_session = Session(cred_provider=final_cred_provider)
 

truefoundry/workflow/example/hello_world_package/workflow.py

@@ -5,9 +5,9 @@ from truefoundry.workflow import PythonTaskConfig, TaskPythonBuild, task, workfl
     task_config=PythonTaskConfig(
         image=TaskPythonBuild(
             python_version="3.9",
-            pip_packages=["truefoundry[workflow]==0.3.0rc7"],
+            pip_packages=["truefoundry[workflow]"],
         ),
-        service_account="tfy-flyte-dataplane-devtest-s3",
+        service_account="tfy-workflows-sa",
     )
 )
 def say_hello() -> str: