truefoundry 0.10.4rc4__py3-none-any.whl → 0.10.5__py3-none-any.whl

truefoundry/_ask/cli.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Tuple
+from typing import Callable, Dict, Optional, Tuple, Union
 
 import rich_click as click
 from openai import AsyncOpenAI
@@ -11,8 +11,6 @@ from truefoundry.cli.const import COMMAND_CLS
 from truefoundry.cli.util import handle_exception_wrapper, select_cluster
 from truefoundry.common.constants import (
     ENV_VARS,
-    OPENAI_API_KEY_KEY,
-    OPENAI_MODEL_KEY,
     TFY_ASK_MODEL_NAME_KEY,
     TFY_ASK_OPENAI_API_KEY_KEY,
     TFY_ASK_OPENAI_BASE_URL_KEY,
@@ -23,77 +21,87 @@ from truefoundry.common.utils import get_tfy_servers_config
 console = Console()
 
 
-def _get_openai_client() -> Tuple[AsyncOpenAI, str]:
+class CustomAsyncOpenAI(AsyncOpenAI):
+    def __init__(
+        self, *, api_key: Optional[Union[str, Callable[[], str]]] = None, **kwargs
+    ):
+        self.__api_key_fn = None
+        if isinstance(api_key, str) or api_key is None:
+            _api_key = api_key
+        else:
+            self.__api_key_fn = api_key
+            _api_key = self.__api_key_fn()
+        super().__init__(api_key=_api_key, **kwargs)
+
+    @property
+    def auth_headers(self) -> Dict[str, str]:
+        if self.__api_key_fn is not None:
+            api_key = self.__api_key_fn()
+        else:
+            api_key = self.api_key
+        return {"Authorization": f"Bearer {api_key}"}
+
+
+def _get_openai_client(session: Session) -> Tuple[CustomAsyncOpenAI, str]:
     """
     Returns an AsyncOpenAI client using either user-provided credentials or TrueFoundry LLM gateway.
     """
-    console.print("")
-    default_model = "gpt-4o"
-    if ENV_VARS.TFY_ASK_OPENAI_BASE_URL and ENV_VARS.TFY_ASK_OPENAI_API_KEY:
+    if ENV_VARS.TFY_ASK_OPENAI_BASE_URL:
         console.print(
-            f"Found custom OpenAI API settings ([green]{TFY_ASK_OPENAI_BASE_URL_KEY}[/green], [green]{TFY_ASK_OPENAI_API_KEY_KEY}[/green]) in env"
+            f"Found custom OpenAI compatible API settings ([green]{TFY_ASK_OPENAI_BASE_URL_KEY}[/green]) in env"
         )
-        client = AsyncOpenAI(
-            base_url=ENV_VARS.TFY_ASK_OPENAI_BASE_URL,
-            api_key=ENV_VARS.TFY_ASK_OPENAI_API_KEY,
-        )
-        if ENV_VARS.TFY_ASK_MODEL_NAME:
-            openai_model = ENV_VARS.TFY_ASK_MODEL_NAME
+        if ENV_VARS.TFY_ASK_OPENAI_API_KEY:
             console.print(
-                f"Using custom OpenAI model from env [green]{TFY_ASK_MODEL_NAME_KEY}[/green]: [yellow]{openai_model}[/yellow]"
+                f"Found API key ([green]{TFY_ASK_OPENAI_API_KEY_KEY}[/green]) in env"
             )
+            api_key = ENV_VARS.TFY_ASK_OPENAI_API_KEY
         else:
-            openai_model = default_model
             console.print(
-                f"Using default OpenAI model: [yellow]{openai_model}[/yellow]"
-                f"\n[dim]Tip: To use a different model, set the env var "
-                f"[green]{TFY_ASK_MODEL_NAME_KEY}[/green] to the model name you want to use.[/dim]"
+                f"No API key found in env, using [yellow]EMPTY[/yellow] as API key"
+                f"\n[dim]Tip: To use a different API key, set the env var "
+                f"[green]{TFY_ASK_OPENAI_API_KEY_KEY}[/green] to the API key you want to use.[/dim]"
             )
+            api_key = "EMPTY"
+        base_url = ENV_VARS.TFY_ASK_OPENAI_BASE_URL
+        default_model = "gpt-4o"
+    else:
+        tfy_servers_config = get_tfy_servers_config(session.tfy_host)
+        base_url = f"{tfy_servers_config.servicefoundry_server_url}/v1/tfy-ai/proxy/api/inference/openai"
         console.print(
-            "[dim][yellow]This operation will use tokens from your model provider and may incur costs.[/yellow][/dim]"
+            f"Using TrueFoundry Managed AI."
+            f"\n[dim]Tip: To use your own OpenAI API compatible API for the ask command, set the following env vars"
+            f"\n * [green]{TFY_ASK_OPENAI_BASE_URL_KEY}[/] to the base URL of your OpenAI compatible API. E.g. [yellow]https://api.openai.com/v1[/yellow]"
+            f"\n * [green]{TFY_ASK_OPENAI_API_KEY_KEY}[/] to the API key of your OpenAI compatible API."
+            f"[/dim]"
         )
         console.print("")
-        return client, openai_model
-    elif ENV_VARS.OPENAI_API_KEY:
-        console.print(f"Found [green]{OPENAI_API_KEY_KEY}[/green] in env")
-        client = AsyncOpenAI(
-            api_key=ENV_VARS.OPENAI_API_KEY,
+
+        api_key = lambda: session.access_token  # noqa: E731
+        default_model = "tfy-ai-openai/gpt-4o"
+    client = CustomAsyncOpenAI(
+        base_url=base_url,
+        api_key=api_key,
+    )
+    if ENV_VARS.TFY_ASK_MODEL_NAME:
+        openai_model = ENV_VARS.TFY_ASK_MODEL_NAME
+        console.print(
+            f"Using custom model from env [green]{TFY_ASK_MODEL_NAME_KEY}[/green]: [yellow]{openai_model}[/yellow]"
         )
-        if ENV_VARS.OPENAI_MODEL:
-            openai_model = ENV_VARS.OPENAI_MODEL
-            console.print(
-                f"Using custom OpenAI model from env [green]{OPENAI_MODEL_KEY}[/green]: [yellow]{openai_model}[/yellow]"
-            )
-        else:
-            openai_model = default_model
-            console.print(
-                f"Using default OpenAI model: [yellow]{openai_model}[/yellow]"
-                f"\n[dim]Tip: To use a different OpenAI model, set the env var "
-                f"[green]{OPENAI_MODEL_KEY}[/green] to the model name you want to use.[/dim]"
-            )
+    else:
+        openai_model = default_model
         console.print(
-            f"[dim]Tip: To use your own OpenAI API compatible model for the ask command, set the env vars "
-            f"[green]{TFY_ASK_OPENAI_BASE_URL_KEY}[/], "
-            f"[green]{TFY_ASK_OPENAI_API_KEY_KEY}[/], and "
-            f"[green]{TFY_ASK_MODEL_NAME_KEY}[/].[/dim]"
+            f"Using default model: [yellow]{openai_model}[/yellow]"
+            f"\n[dim]Tip: To use a different model, set the env var "
+            f"[green]{TFY_ASK_MODEL_NAME_KEY}[/green] to the model name you want to use.[/dim]"
         )
+    console.print("")
+
+    if ENV_VARS.TFY_ASK_OPENAI_BASE_URL:
         console.print(
-            "[dim][yellow]This operation will use tokens from your OpenAI account and may incur costs.[/yellow][/dim]"
+            "[dim][yellow]This operation will use tokens from your model provider and may incur costs.[/yellow][/dim]"
         )
         console.print("")
-        return client, openai_model
-    else:
-        llm_env_instruction = (
-            "No OpenAI API Key found in env."
-            f"\n- To use your own OpenAI API compatible model for the ask command, set the env vars "
-            f"[green]{TFY_ASK_OPENAI_BASE_URL_KEY}[/], "
-            f"[green]{TFY_ASK_OPENAI_API_KEY_KEY}[/], and "
-            f"[green]{TFY_ASK_MODEL_NAME_KEY}[/] (default: {default_model})."
-            f"\n- Alternatively, you can use OpenAI directly by setting the env vars "
-            f"[green]{OPENAI_API_KEY_KEY}[/], "
-            f"[green]{OPENAI_MODEL_KEY}[/] (default: {default_model})"
-        )
-        raise ValueError(llm_env_instruction)
+    return client, openai_model
 
 
 @click.command(name="ask", cls=COMMAND_CLS)
@@ -124,7 +132,7 @@ def ask_command(ctx, cluster: str) -> None:
         "Use this command to ask questions and troubleshoot issues in your Kubernetes cluster managed by the TrueFoundry Control Plane.\n"
         "It helps you investigate and identify potential problems across services, pods, deployments, and more.\n"
     )
-    openai_client, openai_model = _get_openai_client()
+    openai_client, openai_model = _get_openai_client(session=session)
     if not cluster:
         console.print(
             "[dim]Tip: You can specify a cluster using the '--cluster' option, or select one interactively from the list.[/dim]\n"

truefoundry/cli/util.py

@@ -109,7 +109,7 @@ def print_dict_as_table_panel(
 ):
     table = Table(show_header=False, box=None)
     table.add_column("Key", style=f"bold {key_color}", width=15)
-    table.add_column("Value")
+    table.add_column("Value", overflow="fold")
     for key, value in dct.items():
         table.add_row(key, value)
     console.print(

truefoundry/deploy/__init__.py

@@ -118,6 +118,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     SparkImageBuild,
     SparkJobJavaEntrypoint,
     SparkJobPythonEntrypoint,
+    SparkJobPythonNotebookEntrypoint,
     SparkJobScalaEntrypoint,
     SQSInputConfig,
     SQSOutputConfig,

truefoundry/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-06-09T12:01:27+00:00
+#   timestamp: 2025-06-18T21:24:37+00:00
 
 from __future__ import annotations
 
@@ -908,7 +908,7 @@ class SparkImageBuild(BaseModel):
         None,
         description="FQN of the container registry. If you can't find your registry here,\nadd it through the [Integrations](/integrations?tab=docker-registry) page",
     )
-    build_source: GitSource
+    build_source: Union[GitSource, RemoteSource] = Field(..., description="")
     build_spec: SparkBuild
 
 
@@ -936,6 +936,14 @@ class SparkJobPythonEntrypoint(BaseModel):
     )
 
 
+class SparkJobPythonNotebookEntrypoint(BaseModel):
+    type: Literal["python-notebook"] = Field(..., description="")
+    main_application_file: str = Field(
+        ...,
+        description="The main application file to be executed by the spark job. Relative path in case of git repository.",
+    )
+
+
 class SparkJobScalaEntrypoint(BaseModel):
     type: Literal["scala"] = Field(..., description="")
     main_application_file: str = Field(
@@ -1449,7 +1457,10 @@ class SparkJob(BaseModel):
         description="The image to use for driver and executors. Must have spark installed. Spark version must match the version in the image.",
     )
     entrypoint: Union[
-        SparkJobPythonEntrypoint, SparkJobScalaEntrypoint, SparkJobJavaEntrypoint
+        SparkJobPythonEntrypoint,
+        SparkJobScalaEntrypoint,
+        SparkJobJavaEntrypoint,
+        SparkJobPythonNotebookEntrypoint,
     ] = Field(..., description="")
     driver_config: SparkDriverConfig
     executor_config: SparkExecutorConfig

truefoundry/deploy/v2/lib/patched_models.py

@@ -520,3 +520,9 @@ class SparkJobPythonEntrypoint(models.SparkJobPythonEntrypoint, PatchedModelBase
 
 class SparkJobJavaEntrypoint(models.SparkJobJavaEntrypoint, PatchedModelBase):
     type: Literal["java"] = "java"
+
+
+class SparkJobPythonNotebookEntrypoint(
+    models.SparkJobPythonNotebookEntrypoint, PatchedModelBase
+):
+    type: Literal["python-notebook"] = "python-notebook"

{truefoundry-0.10.4rc4.dist-info → truefoundry-0.10.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.10.4rc4
+Version: 0.10.5
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -30,12 +30,12 @@ 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.2.0,>=0.1.1
+Requires-Dist: truefoundry-sdk<0.2.0,>=0.1.3
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0
 Provides-Extra: ai
-Requires-Dist: mcp==1.9.1; (python_version >= '3.10') and extra == 'ai'
+Requires-Dist: mcp==1.9.4; (python_version >= '3.10') and extra == 'ai'
 Provides-Extra: workflow
 Requires-Dist: flytekit==1.15.3; (python_version >= '3.9' and python_version < '3.13') and extra == 'workflow'
 Description-Content-Type: text/markdown

{truefoundry-0.10.4rc4.dist-info → truefoundry-0.10.5.dist-info}/RECORD

@@ -4,7 +4,7 @@ truefoundry/logger.py,sha256=u-YCNjg5HBwE70uQcpjIG64Ghos-K2ulTWaxC03BSj4,714
 truefoundry/pydantic_v1.py,sha256=jSuhGtz0Mbk1qYu8jJ1AcnIDK4oxUsdhALc4spqstmM,345
 truefoundry/version.py,sha256=bqiT4Q-VWrTC6P4qfK43mez-Ppf-smWfrl6DcwV7mrw,137
 truefoundry/_ask/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/_ask/cli.py,sha256=zPaDvxhX2dITmPTtut2Iu6WAIaizrwR-U_dDZ6xv2io,5814
+truefoundry/_ask/cli.py,sha256=RDi1lwbUMYw0CnvaYG4o6o1phmnKjuggdQ5I8sllTlA,5812
 truefoundry/_ask/client.py,sha256=QWQRiDwmtIlLaZsyGcLZaQstYFzpmJeCRdATMapjL-8,18740
 truefoundry/_ask/llm_utils.py,sha256=ayjz7JtVu142lrm8t0cVoxLxUpx76b71y8R62z_WurY,13537
 truefoundry/autodeploy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -37,7 +37,7 @@ truefoundry/cli/config.py,sha256=f7z0_gmYZiNImB7Bxz0AnOlrxY2X4lFnX4jYW1I7NHQ,139
 truefoundry/cli/console.py,sha256=9-dMy4YPisCJQziRKTg8Qa0UJnOGl1soiUnJjsnLDvE,242
 truefoundry/cli/const.py,sha256=dVHPo1uAiDSSMXwXoT2mR5kNQjExT98QNVRz98Hz_Ts,510
 truefoundry/cli/display_util.py,sha256=9vzN3mbQqU6OhS7qRUiMRana4PTHa4sDTA0Hn7OVjCI,3108
-truefoundry/cli/util.py,sha256=pezUfF2GC6ru7s8VeH2a7uvXTU0xN9ka7yLXkIgC3dY,4998
+truefoundry/cli/util.py,sha256=sSKFZ5wGkForoRIiKD2gFlMyj9D1iinzdjQtMYJx8oU,5015
 truefoundry/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/common/auth_service_client.py,sha256=N3YxKlx63r6cPZqbgb2lqBOPI69ShB7D7RCIq4FSCjc,7949
 truefoundry/common/constants.py,sha256=nWd3Je71WmHEORRUTCupZy5fWADqEFftjYP6wiYhCIc,4627
@@ -52,9 +52,9 @@ truefoundry/common/storage_provider_utils.py,sha256=yURhMw8k0FLFvaviRHDiifhvc6Gn
 truefoundry/common/types.py,sha256=BMJFCsR1lPJAw66IQBSvLyV4I6o_x5oj78gVsUa9si8,188
 truefoundry/common/utils.py,sha256=j3QP0uOsaGD_VmDDR68JTwoYE1okkAq6OqpVkzVf48Q,6424
 truefoundry/common/warnings.py,sha256=rs6BHwk7imQYedo07iwh3TWEOywAR3Lqhj0AY4khByg,504
-truefoundry/deploy/__init__.py,sha256=2GNbI8IGJBotz_IKaqQ-DWYWZn_pSu7lN7aId15Gk7Q,2799
+truefoundry/deploy/__init__.py,sha256=PVbGPU9S3-dTFn5LvLwaEnfsp2RrGT9iiM7_15kOV84,2837
 truefoundry/deploy/python_deploy_codegen.py,sha256=k19_m5DGsUyjOUCSKwIVP8vDna2sq01tHABsUfoVpW4,8019
-truefoundry/deploy/_autogen/models.py,sha256=xt-DuaRDx5jeRwyGoQH2yyPZAep9Q2MHFW9XBuRzG8E,73161
+truefoundry/deploy/_autogen/models.py,sha256=8j_y0Yp8k8Sjj7iVtZDHeuxq9kDvD0xI8-iFnbf0370,73571
 truefoundry/deploy/builder/__init__.py,sha256=kgvlkVkiWpMVdim81tIeLrdoACqrFDgwCqHdQVsCsMo,4988
 truefoundry/deploy/builder/constants.py,sha256=amUkHoHvVKzGv0v_knfiioRuKiJM0V0xW0diERgWiI0,508
 truefoundry/deploy/builder/docker_service.py,sha256=sm7GWeIqyrKaZpxskdLejZlsxcZnM3BTDJr6orvPN4E,3948
@@ -117,7 +117,7 @@ truefoundry/deploy/v2/lib/deploy.py,sha256=HfSUdAS3gSpFAFtV0Mq9LscfpkaXqA2LHW4VX
 truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=G5BzMIbap8pgDX1eY-TITruUxQdkKhYtBmRwLL6lDeY,14342
 truefoundry/deploy/v2/lib/deployable_patched_models.py,sha256=mUi-OjPf7bc8rzfrPLdFb79LKuDq7F36RxL4V-AXebs,6830
 truefoundry/deploy/v2/lib/models.py,sha256=ogc1UYs1Z2nBdGSKCrde9sk8d0GxFKMkem99uqO5CmM,1148
-truefoundry/deploy/v2/lib/patched_models.py,sha256=VkfS7akbUzMA4q15lQUcAirdTsyVE1rfMeCmjXJC6Zk,15394
+truefoundry/deploy/v2/lib/patched_models.py,sha256=oNsOr5ojVn2XHjATD3VLuuO6w_ljDL99siHXy6y3Y0g,15558
 truefoundry/deploy/v2/lib/source.py,sha256=d6-8_6Zn5koBglqrBrY6ZLG_7yyPuLdyEmK4iZTw6xY,9405
 truefoundry/ml/__init__.py,sha256=EEEHV7w58Krpo_W9Chd8Y3TdItfFO3LI6j6Izqc4-P8,2219
 truefoundry/ml/constants.py,sha256=vDq72d4C9FSWqr9MMdjgTF4TuyNFApvo_6RVsSeAjB4,2837
@@ -381,7 +381,7 @@ truefoundry/workflow/remote_filesystem/__init__.py,sha256=LQ95ViEjJ7Ts4JcCGOxMPs
 truefoundry/workflow/remote_filesystem/logger.py,sha256=em2l7D6sw7xTLDP0kQSLpgfRRCLpN14Qw85TN7ujQcE,1022
 truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=xcT0wQmQlgzcj0nP3tJopyFSVWT1uv3nhiTIuwfXYeg,12342
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=nSGPZu0Gyd_jz0KsEE-7w_BmnTD8CVF1S8cUJoxaCbc,13305
-truefoundry-0.10.4rc4.dist-info/METADATA,sha256=5AJCpFRmPJFEZb9YerIsMFvC5Dq_Sr3YJPJlh3Ea3IQ,2508
-truefoundry-0.10.4rc4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truefoundry-0.10.4rc4.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
-truefoundry-0.10.4rc4.dist-info/RECORD,,
+truefoundry-0.10.5.dist-info/METADATA,sha256=t8-jiTcLLtOLh7QQQwl_FXqgG0GJdttfh8b06Org1iE,2505
+truefoundry-0.10.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truefoundry-0.10.5.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
+truefoundry-0.10.5.dist-info/RECORD,,