truefoundry 0.7.0rc6__py3-none-any.whl → 0.7.1__py3-none-any.whl

truefoundry/deploy/lib/clients/_mcp_streamable_http.py

@@ -0,0 +1,264 @@
+"""
+StreamableHTTP Client Transport Module
+# From https://github.com/modelcontextprotocol/python-sdk/pull/573
+
+This module implements the StreamableHTTP transport for MCP clients,
+providing support for HTTP POST requests with optional SSE streaming responses
+and session management.
+"""
+
+import logging
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import Any
+
+import anyio
+import httpx
+from httpx_sse import EventSource, aconnect_sse
+from mcp.types import (
+    ErrorData,
+    JSONRPCError,
+    JSONRPCMessage,
+    JSONRPCNotification,
+    JSONRPCRequest,
+)
+
+logger = logging.getLogger(__name__)
+
+# Header names
+MCP_SESSION_ID_HEADER = "mcp-session-id"
+LAST_EVENT_ID_HEADER = "last-event-id"
+
+# Content types
+CONTENT_TYPE_JSON = "application/json"
+CONTENT_TYPE_SSE = "text/event-stream"
+
+
+@asynccontextmanager
+async def streamablehttp_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: timedelta = timedelta(seconds=30),
+    sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+):
+    """
+    Client transport for StreamableHTTP.
+
+    `sse_read_timeout` determines how long (in seconds) the client will wait for a new
+    event before disconnecting. All other HTTP operations are controlled by `timeout`.
+
+    Yields:
+        Tuple of (read_stream, write_stream, terminate_callback)
+    """
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream[
+        JSONRPCMessage | Exception
+    ](0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream[
+        JSONRPCMessage
+    ](0)
+
+    async def get_stream():
+        """
+        Optional GET stream for server-initiated messages
+        """
+        nonlocal session_id
+        try:
+            # Only attempt GET if we have a session ID
+            if not session_id:
+                return
+
+            get_headers = request_headers.copy()
+            get_headers[MCP_SESSION_ID_HEADER] = session_id
+
+            async with aconnect_sse(
+                client,
+                "GET",
+                url,
+                headers=get_headers,
+                timeout=httpx.Timeout(timeout.seconds, read=sse_read_timeout.seconds),
+            ) as event_source:
+                event_source.response.raise_for_status()
+                logger.debug("GET SSE connection established")
+
+                async for sse in event_source.aiter_sse():
+                    if sse.event == "message":
+                        try:
+                            message = JSONRPCMessage.model_validate_json(sse.data)
+                            logger.debug(f"GET message: {message}")
+                            await read_stream_writer.send(message)
+                        except Exception as exc:
+                            logger.error(f"Error parsing GET message: {exc}")
+                            await read_stream_writer.send(exc)
+                    else:
+                        logger.warning(f"Unknown SSE event from GET: {sse.event}")
+        except Exception as exc:
+            # GET stream is optional, so don't propagate errors
+            logger.debug(f"GET stream error (non-fatal): {exc}")
+
+    async def post_writer(client: httpx.AsyncClient):
+        nonlocal session_id
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    # Add session ID to headers if we have one
+                    post_headers = request_headers.copy()
+                    if session_id:
+                        post_headers[MCP_SESSION_ID_HEADER] = session_id
+
+                    logger.debug(f"Sending client message: {message}")
+
+                    # Handle initial initialization request
+                    is_initialization = (
+                        isinstance(message.root, JSONRPCRequest)
+                        and message.root.method == "initialize"
+                    )
+                    if (
+                        isinstance(message.root, JSONRPCNotification)
+                        and message.root.method == "notifications/initialized"
+                    ):
+                        tg.start_soon(get_stream)
+
+                    async with client.stream(
+                        "POST",
+                        url,
+                        json=message.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        ),
+                        headers=post_headers,
+                    ) as response:
+                        if response.status_code == 202:
+                            logger.debug("Received 202 Accepted")
+                            continue
+                        # Check for 404 (session expired/invalid)
+                        if response.status_code == 404:
+                            if isinstance(message.root, JSONRPCRequest):
+                                jsonrpc_error = JSONRPCError(
+                                    jsonrpc="2.0",
+                                    id=message.root.id,
+                                    error=ErrorData(
+                                        code=32600,
+                                        message="Session terminated",
+                                    ),
+                                )
+                                await read_stream_writer.send(
+                                    JSONRPCMessage(jsonrpc_error)
+                                )
+                                continue
+
+                        if not response.is_success:
+                            _response_content = await response.aread()
+                            logger.error(
+                                f"Response: {response.status_code} {_response_content}"
+                            )
+                            response.raise_for_status()
+
+                        # Extract session ID from response headers
+                        if is_initialization:
+                            new_session_id = response.headers.get(MCP_SESSION_ID_HEADER)
+                            if new_session_id:
+                                session_id = new_session_id
+                                logger.info(f"Received session ID: {session_id}")
+
+                        # Handle different response types
+                        content_type = response.headers.get("content-type", "").lower()
+
+                        if content_type.startswith(CONTENT_TYPE_JSON):
+                            try:
+                                content = await response.aread()
+                                json_message = JSONRPCMessage.model_validate_json(
+                                    content
+                                )
+                                await read_stream_writer.send(json_message)
+                            except Exception as exc:
+                                logger.error(f"Error parsing JSON response: {exc}")
+                                await read_stream_writer.send(exc)
+
+                        elif content_type.startswith(CONTENT_TYPE_SSE):
+                            # Parse SSE events from the response
+                            try:
+                                event_source = EventSource(response)
+                                async for sse in event_source.aiter_sse():
+                                    if sse.event == "message":
+                                        try:
+                                            await read_stream_writer.send(
+                                                JSONRPCMessage.model_validate_json(
+                                                    sse.data
+                                                )
+                                            )
+                                        except Exception as exc:
+                                            logger.exception("Error parsing message")
+                                            await read_stream_writer.send(exc)
+                                    else:
+                                        logger.warning(f"Unknown event: {sse.event}")
+
+                            except Exception as e:
+                                logger.exception("Error reading SSE stream:")
+                                await read_stream_writer.send(e)
+
+                        else:
+                            # For 202 Accepted with no body
+                            if response.status_code == 202:
+                                logger.debug("Received 202 Accepted")
+                                continue
+
+                            error_msg = f"Unexpected content type: {content_type}"
+                            logger.error(error_msg)
+                            await read_stream_writer.send(ValueError(error_msg))
+
+        except Exception as exc:
+            logger.error(f"Error in post_writer: {exc}")
+        finally:
+            await read_stream_writer.aclose()
+            await write_stream.aclose()
+
+    async def terminate_session():
+        """
+        Terminate the session by sending a DELETE request.
+        """
+        nonlocal session_id
+        if not session_id:
+            return  # No session to terminate
+
+        try:
+            delete_headers = request_headers.copy()
+            delete_headers[MCP_SESSION_ID_HEADER] = session_id
+
+            response = await client.delete(
+                url,
+                headers=delete_headers,
+            )
+
+            if response.status_code == 405:
+                # Server doesn't allow client-initiated termination
+                logger.debug("Server does not allow session termination")
+            elif response.status_code != 200:
+                logger.warning(f"Session termination failed: {response.status_code}")
+        except Exception as exc:
+            logger.warning(f"Session termination failed: {exc}")
+
+    async with anyio.create_task_group() as tg:
+        try:
+            logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
+            # Set up headers with required Accept header
+            request_headers = {
+                "Accept": f"{CONTENT_TYPE_JSON}, {CONTENT_TYPE_SSE}",
+                "Content-Type": CONTENT_TYPE_JSON,
+                **(headers or {}),
+            }
+            # Track session ID if provided by server
+            session_id: str | None = None
+
+            async with httpx.AsyncClient(
+                headers=request_headers,
+                timeout=httpx.Timeout(timeout.seconds, read=sse_read_timeout.seconds),
+                follow_redirects=True,
+            ) as client:
+                tg.start_soon(post_writer, client)
+                try:
+                    yield read_stream, write_stream, terminate_session
+                finally:
+                    tg.cancel_scope.cancel()
+        finally:
+            await read_stream_writer.aclose()
+            await write_stream.aclose()

truefoundry/deploy/lib/clients/ask_client.py

@@ -0,0 +1,371 @@
+try:
+    from mcp import ClientSession
+    from mcp.types import TextContent
+
+    from truefoundry.deploy.lib.clients._mcp_streamable_http import (
+        streamablehttp_client,
+    )
+except ImportError:
+    import sys
+
+    python_version = sys.version_info
+    raise ImportError(
+        f"This feature requires Python 3.10 or higher. Your current Python version is '{python_version.major}.{python_version.minor}.{python_version.micro}'. "
+        "Please upgrade to a supported version."
+    ) from None
+
+import json
+from contextlib import AsyncExitStack
+from typing import List, Optional, Union
+
+import rich_click as click
+from openai import NOT_GIVEN, AsyncOpenAI
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionToolParam,
+    ChatCompletionUserMessageParam,
+)
+from pydantic import BaseModel
+from rich.console import Console
+from rich.status import Status
+
+from truefoundry.cli.display_util import log_chat_completion_message
+from truefoundry.common.constants import ENV_VARS
+from truefoundry.logger import logger
+
+console = Console(soft_wrap=False)
+
+
+class AskClient:
+    """Handles the chat session lifecycle between the user and the assistant via OpenAI and MCP."""
+
+    def __init__(
+        self,
+        cluster: str,
+        token: str,
+        openai_model: str,
+        debug: bool = False,
+        openai_client: Optional[AsyncOpenAI] = None,
+    ):
+        self.cluster = cluster
+        self.token = token
+        self.debug = debug
+
+        self.async_openai_client = openai_client or AsyncOpenAI()
+        # Initialize the OpenAI client with the session
+        self.openai_model = openai_model
+        self._log_message(
+            f"\nInitialize OpenAI client with model: {self.openai_model!r}, base_url: {str(self.async_openai_client.base_url)!r}\n"
+        )
+
+        self.exit_stack = AsyncExitStack()
+        self.history: List[ChatCompletionMessageParam] = []
+        self._cached_tools: Optional[List[ChatCompletionToolParam]] = None
+
+    async def connect(self, server_url: str, prompt_name: Optional[str] = None):
+        """Initialize connection to the SSE-based MCP server and prepare the chat session."""
+        try:
+            logger.debug(f"Starting a new client for {server_url}")
+            self._streams_context = streamablehttp_client(
+                url=server_url, headers=self._auth_headers()
+            )
+            (
+                read_stream,
+                write_stream,
+                self._terminate_cb,
+            ) = await self._streams_context.__aenter__()
+            self._session_context = ClientSession(
+                read_stream=read_stream, write_stream=write_stream
+            )
+            self.session = await self._session_context.__aenter__()
+            await self.session.initialize()
+            self._log_message("Connected and session initialized.")
+            await self._list_tools()  # Pre-load tool definitions for tool-calling
+            self._log_message(
+                "\nTFY ASK is ready. Type 'exit' to quit.", log=self.debug
+            )
+
+            await self._load_initial_prompt(prompt_name)
+
+        except Exception as e:
+            self._log_message(f"❌ Connection error: {e}")
+            await self.cleanup()
+            raise
+        finally:
+            await self.exit_stack.__aenter__()
+
+    async def cleanup(self):
+        """Properly close all async contexts opened during session initialization."""
+        await self._terminate_cb()
+        for context in [
+            getattr(self, "_session_context", None),
+            getattr(self, "_streams_context", None),
+        ]:
+            if context:
+                await context.__aexit__(None, None, None)
+
+    async def chat_loop(self):
+        """Interactive loop: accepts user queries and returns responses until interrupted or 'exit' is typed."""
+        await self.process_query()  # Optional greeting message from assistant
+
+        while True:
+            try:
+                query = click.prompt(click.style("User", fg="yellow"), type=str)
+                if not query:
+                    self._log_message("Empty query. Type 'exit' to quit.", log=True)
+                    continue
+
+                if query.lower() in ("exit", "quit"):
+                    self._log_message("Exiting chat...")
+                    break
+
+                await self.process_query(query)
+
+            except (KeyboardInterrupt, EOFError, click.Abort):
+                self._log_message("\nChat interrupted.")
+                break
+
+    async def process_query(self, query: Optional[str] = None, max_turns: int = 50):
+        """Handles sending user input to the assistant and processing the assistant’s reply."""
+        if query:
+            self._append_message(
+                ChatCompletionUserMessageParam(role="user", content=query),
+                log=self.debug,
+            )
+
+        tools = await self._list_tools()  # Fetch or use cached tool list
+
+        turn: int = 0
+        # Backup history to revert if OpenAI call fails
+        _checkpoint_idx = len(self.history)
+
+        with console.status(status="Thinking...", spinner="dots") as spinner:
+            while True:
+                try:
+                    if turn >= max_turns:
+                        self._log_message("Max turns reached. Exiting.")
+                        break
+                    spinner.update("Thinking...", spinner="dots")
+                    response = await self._call_openai(
+                        model=self.openai_model, tools=tools
+                    )
+                    turn += 1
+                    message = response.choices[0].message
+
+                    if message.tool_calls:
+                        await self._handle_tool_calls(message, spinner)
+                    elif message.content:
+                        self._append_message(
+                            ChatCompletionAssistantMessageParam(
+                                role="assistant", content=message.content
+                            )
+                        )
+                        break
+                    else:
+                        self._log_message("No assistant response.")
+                        break
+                except Exception as e:
+                    self._log_message(f"OpenAI call failed: {e}", log=self.debug)
+                    console.print(
+                        "Something went wrong. Please try rephrasing your query."
+                    )
+                    self.history = self.history[
+                        :_checkpoint_idx
+                    ]  # Revert to safe state
+                    turn = 0
+                    break
+
+    async def _list_tools(self) -> Optional[List[ChatCompletionToolParam]]:
+        """Fetch and cache the list of available tools from the MCP session."""
+        if self._cached_tools:
+            return self._cached_tools
+
+        self._cached_tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": t.name,
+                    "description": t.description,
+                    "parameters": t.inputSchema,
+                },
+            }
+            for t in (await self.session.list_tools()).tools
+        ]
+
+        self._log_message("\nAvailable tools:")
+        for tool in self._cached_tools or []:
+            self._log_message(
+                f"  - {tool['function']['name']}: {tool['function']['description']}"
+            )
+        return self._cached_tools
+
+    async def _load_initial_prompt(self, prompt_name: Optional[str]) -> None:
+        """Load a system prompt to set assistant behavior at session start."""
+        if not (self.session and prompt_name):
+            return
+
+        result = await self.session.get_prompt(name=prompt_name)
+        if not result:
+            self._log_message("Failed to get initial system prompt.")
+            return
+
+        for message in result.messages:
+            data = message.model_dump() if isinstance(message, BaseModel) else message
+            content = None
+            if isinstance(data, dict):
+                content = (
+                    data.get("content", {}).get("text")
+                    if isinstance(data.get("content"), dict)
+                    else data.get("content")
+                )
+            else:
+                content = data
+
+            # First message is system prompt?
+
+            if content:
+                self._append_message(
+                    ChatCompletionSystemMessageParam(role="system", content=content),
+                    log=self.debug,
+                )
+
+    async def _call_openai(
+        self, model: str, tools: Optional[List[ChatCompletionToolParam]]
+    ):
+        """Make a chat completion request to OpenAI with optional tool support."""
+        return await self.async_openai_client.chat.completions.create(
+            model=model,
+            messages=self.history,
+            tools=tools or NOT_GIVEN,
+            temperature=0.0,  # Set to 0 for deterministic behavior
+            top_p=1,
+        )
+
+    async def _handle_tool_calls(self, message, spinner: Status):
+        """Execute tool calls returned by the assistant and return the results."""
+        for tool_call in message.tool_calls:
+            try:
+                spinner.update(
+                    f"Executing tool: {tool_call.function.name}", spinner="aesthetic"
+                )
+                args = json.loads(tool_call.function.arguments)
+                result = await self.session.call_tool(tool_call.function.name, args)
+                content = getattr(result, "content", result)
+                result_content = self._format_tool_result(content)
+            except Exception as e:
+                result_content = f"Tool `{tool_call.function.name}` call failed: {e}"
+
+            # Log assistant's tool call
+            self._append_message(
+                ChatCompletionAssistantMessageParam(
+                    role="assistant", content=None, tool_calls=[tool_call]
+                ),
+                log=self.debug,
+            )
+
+            # Log tool response
+            self._append_message(
+                ChatCompletionToolMessageParam(
+                    role="tool", tool_call_id=tool_call.id, content=result_content
+                ),
+                log=self.debug,
+            )
+
+    def _format_tool_result(self, content) -> str:
+        """Format tool result into a readable string or JSON block."""
+        if isinstance(content, list):
+            content = (
+                content[0].text
+                if len(content) == 1 and isinstance(content[0], TextContent)
+                else content
+            )
+            if isinstance(content, list):
+                return (
+                    "```\n"
+                    + "\n".join(
+                        (
+                            item.model_dump_json(indent=2)
+                            if isinstance(item, BaseModel)
+                            else str(item)
+                        )
+                        for item in content
+                    )
+                    + "\n```"
+                )
+
+        if isinstance(content, (BaseModel, dict)):
+            return (
+                "```\n"
+                + json.dumps(
+                    content.model_dump() if isinstance(content, BaseModel) else content,
+                    indent=2,
+                )
+                + "\n```"
+            )
+
+        if isinstance(content, str):
+            try:
+                return "```\n" + json.dumps(json.loads(content), indent=2) + "\n```"
+            except Exception:
+                return content
+
+        return str(content)
+
+    def _append_message(self, message: ChatCompletionMessageParam, log: bool = True):
+        """Append a message to history and optionally log it."""
+        self._log_message(message, log)
+        self.history.append(message)
+
+    def _auth_headers(self):
+        """Generate authorization headers for connecting to the SSE server."""
+        return {
+            "Authorization": f"Bearer {self.token}",
+            "X-TFY-Cluster-Id": self.cluster,
+        }
+
+    def _log_message(
+        self,
+        message: Union[str, ChatCompletionMessageParam],
+        log: bool = False,
+    ):
+        """Display a message using Rich console, conditionally based on debug settings."""
+        if not self.debug and not log:
+            return
+        if isinstance(message, str):
+            console.print(message)
+        else:
+            log_chat_completion_message(message, console_=console)
+
+
+async def ask_client_main(
+    cluster: str,
+    server_url: str,
+    token: str,
+    openai_model: str,
+    debug: bool = False,
+    openai_client: Optional[AsyncOpenAI] = None,
+):
+    """Main entrypoint for launching the AskClient chat loop."""
+    ask_client = AskClient(
+        cluster=cluster,
+        token=token,
+        debug=debug,
+        openai_client=openai_client,
+        openai_model=openai_model,
+    )
+    try:
+        await ask_client.connect(
+            server_url=server_url, prompt_name=ENV_VARS.TFY_ASK_SYSTEM_PROMPT_NAME
+        )
+        await ask_client.chat_loop()
+    except Exception as e:
+        console.print(
+            f"[red]An unexpected error occurred while running the assistant: {e}[/red], Check with TrueFoundry support for more details."
+        )
+    except KeyboardInterrupt:
+        console.print("[yellow]Chat interrupted.[/yellow]")
+    finally:
+        await ask_client.cleanup()

truefoundry/workflow/__init__.py

@@ -2,7 +2,10 @@ try:
     import fsspec
     from flytekit import task as _
 except ImportError:
-    print("To use workflows, please run 'pip install truefoundry[workflow]'.")
+    print(
+        "To use workflows, please run 'pip install truefoundry[workflow]'. "
+        "Note: The `workflow` feature is only available for Python 3.9 to 3.12"
+    )
 
 from flytekit import conditional
 from flytekit.types.directory import FlyteDirectory

{truefoundry-0.7.0rc6.dist-info → truefoundry-0.7.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.7.0rc6
+Version: 0.7.1
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -14,9 +14,10 @@ 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.6.0; 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<25.0,>=20.0
+Requires-Dist: packaging<26.0,>=20.0
 Requires-Dist: pydantic<3.0.0,>=1.8.2
 Requires-Dist: pygments<3.0.0,>=2.12.0
 Requires-Dist: pyjwt<3.0.0,>=2.0.0
@@ -30,7 +31,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.15
+Requires-Dist: truefoundry-sdk==0.0.16
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0