truefoundry 0.9.1__py3-none-any.whl → 0.9.2rc1__py3-none-any.whl

truefoundry/_ask/llm_utils.py

@@ -0,0 +1,344 @@
+import copy
+import datetime
+import json
+import re
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, cast
+
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionMessageToolCall,
+    ChatCompletionMessageToolCallParam,
+    ChatCompletionToolParam,
+)
+from rich.console import Console, Group
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.text import Text
+
+# Taken from https://github.com/pydantic/pydantic-ai/blob/222bec41e28fb96b49e71609c43b683d97d1dc97/pydantic_ai_slim/pydantic_ai/models/gemini.py#L791
+
+JsonSchema = dict[str, Any]
+
+
+@dataclass(init=False)
+class WalkJsonSchema(ABC):
+    """Walks a JSON schema, applying transformations to it at each level.
+
+    Note: We may eventually want to rework tools to build the JSON schema from the type directly, using a subclass of
+    pydantic.json_schema.GenerateJsonSchema, rather than making use of this machinery.
+    """
+
+    def __init__(
+        self,
+        schema: JsonSchema,
+        *,
+        prefer_inlined_defs: bool = False,
+        simplify_nullable_unions: bool = False,
+    ):
+        self.schema = schema
+        self.prefer_inlined_defs = prefer_inlined_defs
+        self.simplify_nullable_unions = simplify_nullable_unions
+
+        self.defs: dict[str, JsonSchema] = self.schema.get("$defs", {})
+        self.refs_stack: list[str] = []
+        self.recursive_refs = set[str]()
+
+    @abstractmethod
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        """Make changes to the schema."""
+        return schema
+
+    def walk(self) -> JsonSchema:
+        schema = copy.deepcopy(self.schema)
+
+        # First, handle everything but $defs:
+        schema.pop("$defs", None)
+        handled = self._handle(schema)
+
+        if not self.prefer_inlined_defs and self.defs:
+            handled["$defs"] = {k: self._handle(v) for k, v in self.defs.items()}
+
+        elif self.recursive_refs:  # pragma: no cover
+            # If we are preferring inlined defs and there are recursive refs, we _have_ to use a $defs+$ref structure
+            # We try to use whatever the original root key was, but if it is already in use,
+            # we modify it to avoid collisions.
+            defs = {key: self.defs[key] for key in self.recursive_refs}
+            root_ref = self.schema.get("$ref")
+            root_key = None if root_ref is None else re.sub(r"^#/\$defs/", "", root_ref)
+            if root_key is None:
+                root_key = self.schema.get("title", "root")
+                while root_key in defs:
+                    # Modify the root key until it is not already in use
+                    root_key = f"{root_key}_root"
+
+            defs[root_key] = handled
+            return {"$defs": defs, "$ref": f"#/$defs/{root_key}"}
+
+        return handled
+
+    def _handle(self, schema: JsonSchema) -> JsonSchema:
+        nested_refs = 0
+        if self.prefer_inlined_defs:
+            while ref := schema.get("$ref"):
+                key = re.sub(r"^#/\$defs/", "", ref)
+                if key in self.refs_stack:
+                    self.recursive_refs.add(key)
+                    break  # recursive ref can't be unpacked
+                self.refs_stack.append(key)
+                nested_refs += 1
+
+                def_schema = self.defs.get(key)
+                if def_schema is None:  # pragma: no cover
+                    raise ValueError(f"Could not find $ref definition for {key}")
+                schema = def_schema
+
+        # Handle the schema based on its type / structure
+        type_ = schema.get("type")
+        if type_ == "object":
+            schema = self._handle_object(schema)
+        elif type_ == "array":
+            schema = self._handle_array(schema)
+        elif type_ is None:
+            schema = self._handle_union(schema, "anyOf")
+            schema = self._handle_union(schema, "oneOf")
+
+        # Apply the base transform
+        schema = self.transform(schema)
+
+        if nested_refs > 0:
+            self.refs_stack = self.refs_stack[:-nested_refs]
+
+        return schema
+
+    def _handle_object(self, schema: JsonSchema) -> JsonSchema:
+        if properties := schema.get("properties"):
+            handled_properties = {}
+            for key, value in properties.items():
+                handled_properties[key] = self._handle(value)
+            schema["properties"] = handled_properties
+
+        if (additional_properties := schema.get("additionalProperties")) is not None:
+            if isinstance(additional_properties, bool):
+                schema["additionalProperties"] = additional_properties
+            else:
+                schema["additionalProperties"] = self._handle(additional_properties)
+
+        if (pattern_properties := schema.get("patternProperties")) is not None:
+            handled_pattern_properties = {}
+            for key, value in pattern_properties.items():
+                handled_pattern_properties[key] = self._handle(value)
+            schema["patternProperties"] = handled_pattern_properties
+
+        return schema
+
+    def _handle_array(self, schema: JsonSchema) -> JsonSchema:
+        if prefix_items := schema.get("prefixItems"):
+            schema["prefixItems"] = [self._handle(item) for item in prefix_items]
+
+        if items := schema.get("items"):
+            schema["items"] = self._handle(items)
+
+        return schema
+
+    def _handle_union(
+        self, schema: JsonSchema, union_kind: Literal["anyOf", "oneOf"]
+    ) -> JsonSchema:
+        members = schema.get(union_kind)
+        if not members:
+            return schema
+
+        handled = [self._handle(member) for member in members]
+
+        # convert nullable unions to nullable types
+        if self.simplify_nullable_unions:
+            handled = self._simplify_nullable_union(handled)
+
+        if len(handled) == 1:
+            # In this case, no need to retain the union
+            return handled[0]
+
+        # If we have keys besides the union kind (such as title or discriminator), keep them without modifications
+        schema = schema.copy()
+        schema[union_kind] = handled
+        return schema
+
+    @staticmethod
+    def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
+        # TODO: Should we move this to relevant subclasses? Or is it worth keeping here to make reuse easier?
+        if len(cases) == 2 and {"type": "null"} in cases:
+            # Find the non-null schema
+            non_null_schema = next(
+                (item for item in cases if item != {"type": "null"}),
+                None,
+            )
+            if non_null_schema:
+                # Create a new schema based on the non-null part, mark as nullable
+                new_schema = copy.deepcopy(non_null_schema)
+                new_schema["nullable"] = True
+                return [new_schema]
+            else:  # pragma: no cover
+                # they are both null, so just return one of them
+                return [cases[0]]
+
+        return cases  # pragma: no cover
+
+
+class _GeminiJsonSchema(WalkJsonSchema):
+    """Transforms the JSON Schema from Pydantic to be suitable for Gemini.
+
+    Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
+    a subset of OpenAPI v3.0.3.
+
+    Specifically:
+    * gemini doesn't allow the `title` keyword to be set
+    * gemini doesn't allow `$defs` — we need to inline the definitions where possible
+    """
+
+    def __init__(self, schema: JsonSchema):
+        super().__init__(
+            schema, prefer_inlined_defs=True, simplify_nullable_unions=True
+        )
+
+    def transform(self, schema: JsonSchema) -> JsonSchema:  # noqa: C901
+        # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
+        additional_properties = schema.pop(
+            "additionalProperties", None
+        )  # don't pop yet so it's included in the warning
+        if additional_properties:
+            original_schema = {**schema, "additionalProperties": additional_properties}
+            warnings.warn(
+                "`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema."
+                f" Full schema: {self.schema}\n\n"
+                f"Source of additionalProperties within the full schema: {original_schema}\n\n"
+                "If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n"
+                "If Google's APIs are updated to support this properly, please create an issue on the PydanticAI GitHub"
+                " and we will fix this behavior.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        schema.pop("title", None)
+        schema.pop("default", None)
+        schema.pop("$schema", None)
+        if (const := schema.pop("const", None)) is not None:  # pragma: no cover
+            # Gemini doesn't support const, but it does support enum with a single value
+            schema["enum"] = [const]
+        schema.pop("discriminator", None)
+        schema.pop("examples", None)
+
+        # TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
+        #   where we add notes about these properties to the field description?
+        schema.pop("exclusiveMaximum", None)
+        schema.pop("exclusiveMinimum", None)
+
+        # Gemini only supports string enums, so we need to convert any enum values to strings.
+        # Pydantic will take care of transforming the transformed string values to the correct type.
+        if enum := schema.get("enum"):
+            schema["type"] = "string"
+            schema["enum"] = [str(val) for val in enum]
+
+        type_ = schema.get("type")
+        if "oneOf" in schema and "type" not in schema:  # pragma: no cover
+            # This gets hit when we have a discriminated union
+            # Gemini returns an API error in this case even though it says in its error message it shouldn't...
+            # Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
+            schema["anyOf"] = schema.pop("oneOf")
+
+        if type_ == "string" and (fmt := schema.pop("format", None)):
+            description = schema.get("description")
+            if description:
+                schema["description"] = f"{description} (format: {fmt})"
+            else:
+                schema["description"] = f"Format: {fmt}"
+
+        if "$ref" in schema:
+            raise ValueError(
+                f"Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema['$ref']}"
+            )
+
+        if "prefixItems" in schema:
+            # prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
+            prefix_items = schema.pop("prefixItems")
+            items = schema.get("items")
+            unique_items = [items] if items is not None else []
+            for item in prefix_items:
+                if item not in unique_items:
+                    unique_items.append(item)
+            if len(unique_items) > 1:  # pragma: no cover
+                schema["items"] = {"anyOf": unique_items}
+            elif len(unique_items) == 1:
+                schema["items"] = unique_items[0]
+            schema.setdefault("minItems", len(prefix_items))
+            if items is None:
+                schema.setdefault("maxItems", len(prefix_items))
+
+        return schema
+
+
+def translate_tools_for_gemini(
+    tools: Optional[List[ChatCompletionToolParam]],
+) -> Optional[List[ChatCompletionToolParam]]:
+    if tools is None:
+        return None
+    tools = copy.deepcopy(tools)
+    for tool in tools:
+        if "parameters" in tool["function"]:
+            tool["function"]["parameters"] = _GeminiJsonSchema(
+                tool["function"]["parameters"]
+            ).walk()
+
+    return tools
+
+
+def log_chat_completion_message(
+    message: ChatCompletionMessageParam, console: Console
+) -> None:
+    timestamp = datetime.datetime.now().astimezone().strftime("%Y-%m-%d %H:%M:%S")
+    code_theme = "github-dark"
+    try:
+        _content = message.get("content") or ""
+        message_content = f"```json\n{json.dumps(json.loads(_content), indent=2)}\n```"
+    except (TypeError, ValueError):
+        message_content = str(message.get("content") or "")
+
+    tool_calls: Iterable[ChatCompletionMessageToolCallParam] = cast(
+        list[ChatCompletionMessageToolCallParam],
+        message.get("tool_calls", []),
+    )
+
+    rendered_content: list[Markdown | Text] = []
+
+    if bool(message_content.strip()):
+        rendered_content.append(Markdown(markup=message_content, code_theme=code_theme))
+
+    for call in tool_calls:
+        assert isinstance(call, ChatCompletionMessageToolCall)
+        name = call.function.name
+        args = call.function.arguments
+        rendered_content.append(
+            Text.from_markup(
+                "[bold magenta]Tool Calls:[/bold magenta]\n",
+                overflow="fold",
+            )
+        )
+        rendered_content.append(
+            Markdown(markup=f"```python\n▶ {name}({args})\n```", code_theme=code_theme)
+        )
+
+    if not rendered_content:
+        return
+
+    panel = Panel(
+        Group(*rendered_content, fit=True),
+        title=f"[bold blue]{timestamp}[/bold blue]",
+        title_align="left",
+        border_style="bright_blue",
+        padding=(1, 2),
+        expand=True,
+        width=console.width,
+    )
+    console.print(panel)

truefoundry/cli/__main__.py

@@ -4,6 +4,7 @@ import sys
 import rich_click as click
 
 from truefoundry import logger
+from truefoundry._ask.cli import get_ask_command
 from truefoundry.cli.config import CliConfig
 from truefoundry.cli.const import GROUP_CLS
 from truefoundry.cli.util import setup_rich_click
@@ -11,7 +12,6 @@ from truefoundry.common.constants import TFY_DEBUG_ENV_KEY
 from truefoundry.common.utils import is_internal_env_set
 from truefoundry.deploy.cli.commands import (
     get_apply_command,
-    get_ask_command,
     get_build_command,
     get_delete_command,
     get_deploy_command,

truefoundry/cli/display_util.py

@@ -1,20 +1,9 @@
 import datetime
 import json
-from collections.abc import Iterable
-from typing import Optional, cast
-
-from openai.types.chat import (
-    ChatCompletionMessageParam,
-    ChatCompletionMessageToolCall,
-    ChatCompletionMessageToolCallParam,
-)
+
 from rich import box
 from rich import print_json as _rich_print_json
-from rich.console import Console, Group
-from rich.markdown import Markdown
-from rich.panel import Panel
 from rich.table import Table
-from rich.text import Text
 
 from truefoundry.cli.config import CliConfig
 from truefoundry.cli.console import console
@@ -132,54 +121,3 @@ def print_entity_obj(title, entity):
     for key in keys:
         table.add_row(f"[bold]{stringify(key)}[/]", stringify(entity_data[key]))
     console.print(table)
-
-
-def log_chat_completion_message(
-    message: ChatCompletionMessageParam, console_: Optional[Console] = None
-) -> None:
-    timestamp = datetime.datetime.now().astimezone().strftime("%Y-%m-%d %H:%M:%S")
-    code_theme = "github-dark"
-    target_console = console_ or console
-    try:
-        _content = message.get("content") or ""
-        message_content = f"```json\n{json.dumps(json.loads(_content), indent=2)}\n```"
-    except (TypeError, ValueError):
-        message_content = str(message.get("content") or "")
-
-    tool_calls: Iterable[ChatCompletionMessageToolCallParam] = cast(
-        list[ChatCompletionMessageToolCallParam],
-        message.get("tool_calls", []),
-    )
-
-    rendered_content: list[Markdown | Text] = []
-
-    if bool(message_content.strip()):
-        rendered_content.append(Markdown(markup=message_content, code_theme=code_theme))
-
-    for call in tool_calls:
-        assert isinstance(call, ChatCompletionMessageToolCall)
-        name = call.function.name
-        args = call.function.arguments
-        rendered_content.append(
-            Text.from_markup(
-                "[bold magenta]Tool Calls:[/bold magenta]\n",
-                overflow="fold",
-            )
-        )
-        rendered_content.append(
-            Markdown(markup=f"```python\n▶ {name}({args})\n```", code_theme=code_theme)
-        )
-
-    if not rendered_content:
-        return
-
-    panel = Panel(
-        Group(*rendered_content, fit=True),
-        title=f"[bold blue]{timestamp}[/bold blue]",
-        title_align="left",
-        border_style="bright_blue",
-        padding=(1, 2),
-        expand=True,
-        width=console.width,
-    )
-    target_console.print(panel)

truefoundry/cli/util.py

@@ -1,8 +1,9 @@
 import functools
 import sys
 import zipfile
-from typing import Dict
+from typing import Dict, Optional
 
+import questionary
 import rich_click as click
 from packaging.version import parse as parse_version
 from requests.exceptions import ConnectionError, Timeout
@@ -135,3 +136,33 @@ def _prompt_if_no_value_and_supported(prompt: str, hide_input: bool = True):
         kwargs = {"prompt": prompt, "hide_input": hide_input, "prompt_required": False}
 
     return kwargs
+
+
+def select_cluster(cluster: Optional[str] = None) -> str:
+    """
+    Retrieve available clusters and either return the specified one after validation
+    or allow the user to interactively select from the list.
+    """
+    from truefoundry.deploy.lib.clients.servicefoundry_client import (
+        ServiceFoundryServiceClient,
+    )
+
+    clusters = ServiceFoundryServiceClient().list_clusters()
+
+    if not clusters:
+        raise click.ClickException("No clusters found in your account.")
+
+    if cluster:
+        if not any(c.id == cluster for c in clusters):
+            raise click.ClickException(
+                f"Cluster {cluster} not found. Either it does not exist or you might not be authorized to access it"
+            )
+        return cluster
+
+    choices = {cluster.id: cluster for cluster in clusters}
+    selected_cluster = questionary.select(
+        "Pick a Cluster:", choices=list(choices.keys())
+    ).ask()
+    if not selected_cluster:
+        raise click.ClickException("No cluster selected.")
+    return selected_cluster

truefoundry/common/constants.py

@@ -1,7 +1,7 @@
 import enum
 import os
 from pathlib import Path
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from truefoundry.pydantic_v1 import BaseSettings, Field, SecretStr
 
@@ -26,8 +26,6 @@ TFY_ASK_OPENAI_API_KEY_KEY = "TFY_ASK_OPENAI_API_KEY"
 TFY_ASK_OPENAI_BASE_URL_KEY = "TFY_ASK_OPENAI_BASE_URL"
 TFY_ASK_MODEL_NAME_KEY = "TFY_ASK_OPENAI_MODEL"
 
-TFY_ASK_SYSTEM_PROMPT_NAME_KEY = "TFY_ASK_SYSTEM_PROMPT_NAME"
-
 
 class PythonPackageManager(str, enum.Enum):
     PIP = "pip"
@@ -85,10 +83,11 @@ class TrueFoundrySdkEnv(BaseSettings):
         default=None, env=TFY_ASK_OPENAI_BASE_URL_KEY
     )
     TFY_ASK_MODEL_NAME: Optional[str] = Field(default=None, env=TFY_ASK_MODEL_NAME_KEY)
-
-    TFY_ASK_SYSTEM_PROMPT_NAME: str = Field(
-        default="tfy-ask-k8s-prompt", env=TFY_ASK_SYSTEM_PROMPT_NAME_KEY
+    TFY_ASK_GENERATION_PARAMS: Dict[str, Any] = Field(
+        default_factory=lambda: {"temperature": 0.0, "top_p": 1, "max_tokens": 4096}
     )
+    TFY_ASK_SYSTEM_PROMPT_NAME: str = Field(default="tfy-ask-k8s-prompt")
+    TFY_INTERNAL_ASK_CONFIG_OVERRIDE_FILE: Optional[str] = Field(default=None)
 
     # This is a hack to fresh read the env vars because people can end up importing this file
     # before setting the correct env vars. E.g. in notebook environments.

truefoundry/deploy/cli/commands/__init__.py

@@ -1,5 +1,4 @@
 from truefoundry.deploy.cli.commands.apply_command import get_apply_command
-from truefoundry.deploy.cli.commands.ask_command import get_ask_command
 from truefoundry.deploy.cli.commands.build_command import get_build_command
 from truefoundry.deploy.cli.commands.delete_command import get_delete_command
 from truefoundry.deploy.cli.commands.deploy_command import get_deploy_command
@@ -14,3 +13,19 @@ from truefoundry.deploy.cli.commands.patch_application_command import (
 from truefoundry.deploy.cli.commands.patch_command import get_patch_command
 from truefoundry.deploy.cli.commands.terminate_comand import get_terminate_command
 from truefoundry.deploy.cli.commands.trigger_command import get_trigger_command
+
+__all__ = [
+    "get_apply_command",
+    "get_build_command",
+    "get_delete_command",
+    "get_deploy_command",
+    "get_deploy_init_command",
+    "get_get_command",
+    "get_login_command",
+    "get_logout_command",
+    "get_logs_command",
+    "get_patch_application_command",
+    "get_patch_command",
+    "get_terminate_command",
+    "get_trigger_command",
+]

truefoundry/deploy/cli/commands/kubeconfig_command.py

@@ -5,7 +5,7 @@ import rich_click as click
 from rich.console import Console
 
 from truefoundry.cli.const import COMMAND_CLS
-from truefoundry.cli.util import handle_exception_wrapper
+from truefoundry.cli.util import handle_exception_wrapper, select_cluster
 from truefoundry.common.session import Session
 from truefoundry.deploy.cli.commands.utils import (
     CONTEXT_NAME_FORMAT,
@@ -14,7 +14,6 @@ from truefoundry.deploy.cli.commands.utils import (
     get_kubeconfig_content,
     get_kubeconfig_path,
     save_kubeconfig,
-    select_cluster,
 )
 
 console = Console()

truefoundry/deploy/cli/commands/utils.py

@@ -2,8 +2,6 @@ import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
-import questionary
-import rich_click as click
 import yaml
 
 DEFAULT_KUBECONFIG_PATH: Path = Path.home() / ".kube" / "config"
@@ -114,35 +112,7 @@ def get_cluster_context(
 
 
 def get_cluster_server_url(config: Dict[str, Any], cluster: str) -> Optional[str]:
-    cluster: Optional[Dict[str, Any]] = get_cluster_context(config, cluster)
-    if cluster:
-        return cluster["cluster"].get("server")
+    cluster_context: Optional[Dict[str, Any]] = get_cluster_context(config, cluster)
+    if cluster_context:
+        return cluster_context["cluster"].get("server")
     return None
-
-
-def select_cluster(cluster: Optional[str] = None) -> str:
-    """
-    Retrieve available clusters and either return the specified one after validation
-    or allow the user to interactively select from the list.
-    """
-    from truefoundry.deploy.lib.clients.servicefoundry_client import (
-        ServiceFoundryServiceClient,
-    )
-
-    clusters = ServiceFoundryServiceClient().list_clusters()
-
-    if not clusters:
-        raise click.ClickException("No clusters found in your account.")
-
-    if cluster:
-        if not any(c.id == cluster for c in clusters):
-            raise click.ClickException(
-                f"Cluster {cluster} not found. Either it does not exist or you might not be autthorized to access it"
-            )
-        return cluster
-
-    choices = {cluster.id: cluster for cluster in clusters}
-    cluster = questionary.select("Pick a Cluster:", choices=list(choices.keys())).ask()
-    if not cluster:
-        raise click.ClickException("No cluster selected.")
-    return cluster

truefoundry/deploy/v2/lib/deployable_patched_models.py

@@ -100,10 +100,9 @@ class Application(models.Application, DeployablePatchedModelBase):
     # Unfortunately cue export cannot add discriminator in OAS
     # Even if we add it manually in OAS, `datamodel-code-generator` has bugs when discriminator field is enum type in member models.
     # It will change the members to be incorrect like this
-    # class Service(BaseModel):
-    #     type: Literal["Service"] = Field("Service")  # notice the capital casing
+    # >>> class Service(BaseModel):
+    # >>>     type: Literal["Service"] = Field("Service")  # notice the capital casing
     # This is why we add it manually here
-
     __root__: Union[
         models.Service,
         models.AsyncService,

{truefoundry-0.9.1.dist-info → truefoundry-0.9.2rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.9.1
+Version: 0.9.2rc1
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -14,7 +14,6 @@ Requires-Dist: gitpython<4.0.0,>=3.1.43
 Requires-Dist: importlib-metadata<9.0.0,>=4.11.3
 Requires-Dist: importlib-resources<7.0.0,>=5.2.0
 Requires-Dist: mako<2.0.0,>=1.1.6
-Requires-Dist: mcp==1.8.1; python_version >= '3.10'
 Requires-Dist: numpy<3.0.0,>=1.23.0
 Requires-Dist: openai<2.0.0,>=1.16.2
 Requires-Dist: packaging<26.0,>=20.0
@@ -35,6 +34,8 @@ Requires-Dist: truefoundry-sdk<0.2.0,>=0.1.1
 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'
 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