clarifai 11.4.2__py3-none-any.whl → 11.4.3rc1__py3-none-any.whl

clarifai/runners/models/mcp_class.py

@@ -0,0 +1,143 @@
+"""Base class for creating Model Context Protocol (MCP) servers."""
+
+import asyncio
+import json
+from typing import Any
+
+from fastmcp import Client, FastMCP  # use fastmcp v2 not the built in mcp
+from mcp import types
+from mcp.shared.exceptions import McpError
+
+from clarifai.runners.models.model_class import ModelClass
+
+# class MCPServerProvider(abc.ABC):
+#     """
+#     Base class for creating Model Context Protocol (MCP) servers.
+
+#     This class provides a base implementation of the MCP server, including
+#     methods for handling requests and responses, as well as error handling and
+#     logging.
+
+#     Attributes:
+#     _server: The FastMCP server instance.
+#     _tools: List of tools available in the server.
+#     _resources: List of resources available in the server.
+#     _prompts: List of prompts available in the server.
+
+#     Methods:
+#     get_server(): Returns the FastMCP server instance.
+#     mcp_transport(msg): Handles incoming messages and sends them to the FastMCP server.
+#     """
+
+#     @abc.abstractmethod
+#     def get_server(self) -> FastMCP:
+#         """Required method for each subclass to implement to return the FastMCP server to use."""
+#         if self._server is None:
+#             raise ValueError("Server not initialized")
+#         return self._server
+
+
+class MCPModelClass(ModelClass):
+    """Base class for wrapping FastMCP servers as a model running in Clarfai. This handles
+    all the transport between the API and the MCP server here. Simply subclass this and implement
+    the get_server() method to return the FastMCP server instance. The server is then used to
+    handle all the requests and responses.
+    """
+
+    def load_model(self):
+        # in memory transport provided in fastmcp v2 so we can easily use the client functions.
+        self.client = Client(self.get_server())
+
+    def get_server(self) -> FastMCP:
+        """Required method for each subclass to implement to return the FastMCP server to use."""
+        raise NotImplementedError("Subclasses must implement get_server() method")
+
+    @ModelClass.method
+    def mcp_transport(self, msg: str) -> str:
+        """The single model method to get the jsonrpc message and send it to the FastMCP server then
+        return it's response.
+
+        """
+
+        async def send_notification(client_message: types.ClientNotification) -> None:
+            async with self.client:
+                # Strip the jsonrpc field since send_notification will also pass it in for some reason.
+                client_message = types.ClientNotification.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc"}
+                    )
+                )
+                try:
+                    return await self.client.session.send_notification(client_message)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", error=e.error)
+
+        async def send_request(client_message: types.ClientRequest, id: str) -> Any:
+            async with self.client:
+                # Strip the jsonrpc and id fields as send_request sets them again too.
+                client_message = types.ClientRequest.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc", "id"}
+                    )
+                )
+
+                result_type = None
+                print("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
+                print(types.PingRequest)
+                if isinstance(client_message.root, types.PingRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.InitializeRequest):
+                    return await self.client.session.initialize()
+                elif isinstance(client_message.root, types.SetLevelRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListResourcesRequest):
+                    result_type = types.ListResourcesResult
+                elif isinstance(client_message.root, types.ListResourceTemplatesRequest):
+                    result_type = types.ListResourceTemplatesResult
+                elif isinstance(client_message.root, types.ReadResourceRequest):
+                    result_type = types.ReadResourceResult
+                elif isinstance(client_message.root, types.SubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.UnsubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListPromptsRequest):
+                    result_type = types.ListPromptsResult
+                elif isinstance(client_message.root, types.GetPromptRequest):
+                    result_type = types.GetPromptResult
+                elif isinstance(client_message.root, types.CompleteRequest):
+                    result_type = types.CompleteResult
+                elif isinstance(client_message.root, types.ListToolsRequest):
+                    result_type = types.ListToolsResult
+                elif isinstance(client_message.root, types.CallToolRequest):
+                    result_type = types.CallToolResult
+                else:
+                    # this is a special case where we need to return the list of tools.
+                    raise NotImplementedError(f"Method {client_message.method} not implemented")
+                # Call the mcp server using send_request() or send_notification() depending on the method.
+                try:
+                    return await self.client.session.send_request(client_message, result_type)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", id=id, error=e.error)
+
+        # The message coming here is the generic request. We look at it's .method
+        # to determine which client function to call and to further subparse the params.
+        # Note(zeiler): unfortunately the pydantic types in mcp/types.py are not consistent.
+        # The JSONRPCRequest are supposed to have an id but the InitializeRequest
+        # does not have it.
+        d = json.loads(msg)
+
+        # If we have an id it's a JSONRPCRequest
+        if not d.get('method', None).startswith("notifications/"):
+            # rpc_message2 = types.JSONRPCRequest.model_validate(rpc_message)
+            # underlying: types.JSONRPCRequest = rpc_message.root
+            client_message = types.ClientRequest.model_validate(d)
+            response = asyncio.run(
+                send_request(client_message, id=d.get('id', ""))
+            )  # underlying.id))
+        else:  # JSONRPCRequest
+            client_message = types.ClientNotification.model_validate(d)
+            response = asyncio.run(send_notification(client_message))
+        if response is None:
+            return ""
+        # return as a serialized json string
+        return response.model_dump_json(by_alias=True, exclude_none=True)

clarifai/runners/models/mcp_class.py~

@@ -0,0 +1,149 @@
+"""Base class for creating Model Context Protocol (MCP) servers."""
+
+import asyncio
+import json
+from typing import Any
+
+from fastmcp import Client, FastMCP  # use fastmcp v2 not the built in mcp
+from mcp import types
+from mcp.shared.exceptions import McpError
+
+from clarifai.runners.models.model_class import ModelClass
+
+# class MCPServerProvider(abc.ABC):
+#     """
+#     Base class for creating Model Context Protocol (MCP) servers.
+
+#     This class provides a base implementation of the MCP server, including
+#     methods for handling requests and responses, as well as error handling and
+#     logging.
+
+#     Attributes:
+#     _server: The FastMCP server instance.
+#     _tools: List of tools available in the server.
+#     _resources: List of resources available in the server.
+#     _prompts: List of prompts available in the server.
+
+#     Methods:
+#     get_server(): Returns the FastMCP server instance.
+#     mcp_transport(msg): Handles incoming messages and sends them to the FastMCP server.
+#     """
+
+#     @abc.abstractmethod
+#     def get_server(self) -> FastMCP:
+#         """Required method for each subclass to implement to return the FastMCP server to use."""
+#         if self._server is None:
+#             raise ValueError("Server not initialized")
+#         return self._server
+
+
+class MCPModelClass(ModelClass, MCPServerProvider):
+    """Base class for wrapping FastMCP servers as a model running in Clarfai. This handles
+    all the transport between the API and the MCP server here. Simply subclass this and implement
+    the get_server() method to return the FastMCP server instance. The server is then used to
+    handle all the requests and responses.
+    """
+
+    def load_model(self):
+        # in memory transport provided in fastmcp v2 so we can easily use the client functions.
+        self.client = Client(self.get_server())
+
+    def get_server(self) -> FastMCP:
+        """Required method for each subclass to implement to return the FastMCP server to use."""
+        if self._server is None:
+            raise ValueError("Server not initialized")
+        return self._server
+
+    @ModelClass.method
+    def mcp_transport(self, msg: str) -> str:
+        """The single model method to get the jsonrpc message and send it to the FastMCP server then
+        return it's response.
+
+        Arguments:
+          msg: The incoming message to be handled in serialized JSONRPC format from an MCP client.
+        Returns:
+          str: The response to the incoming message in serialized JSONRPC format
+        """
+
+        async def send_notification(client_message: types.ClientNotification) -> None:
+            async with self.client:
+                # Strip the jsonrpc field since send_notification will also pass it in for some reason.
+                client_message = types.ClientNotification.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc"}
+                    )
+                )
+                try:
+                    return await self.client.session.send_notification(client_message)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", error=e.error)
+
+        async def send_request(client_message: types.ClientRequest, id: str) -> Any:
+            async with self.client:
+                # Strip the jsonrpc and id fields as send_request sets them again too.
+                client_message = types.ClientRequest.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc", "id"}
+                    )
+                )
+
+                result_type = None
+                print("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
+                print(types.PingRequest)
+                if isinstance(client_message.root, types.PingRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.InitializeRequest):
+                    return await self.client.session.initialize()
+                elif isinstance(client_message.root, types.SetLevelRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListResourcesRequest):
+                    result_type = types.ListResourcesResult
+                elif isinstance(client_message.root, types.ListResourceTemplatesRequest):
+                    result_type = types.ListResourceTemplatesResult
+                elif isinstance(client_message.root, types.ReadResourceRequest):
+                    result_type = types.ReadResourceResult
+                elif isinstance(client_message.root, types.SubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.UnsubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListPromptsRequest):
+                    result_type = types.ListPromptsResult
+                elif isinstance(client_message.root, types.GetPromptRequest):
+                    result_type = types.GetPromptResult
+                elif isinstance(client_message.root, types.CompleteRequest):
+                    result_type = types.CompleteResult
+                elif isinstance(client_message.root, types.ListToolsRequest):
+                    result_type = types.ListToolsResult
+                elif isinstance(client_message.root, types.CallToolRequest):
+                    result_type = types.CallToolResult
+                else:
+                    # this is a special case where we need to return the list of tools.
+                    raise NotImplementedError(f"Method {client_message.method} not implemented")
+                # Call the mcp server using send_request() or send_notification() depending on the method.
+                try:
+                    return await self.client.session.send_request(client_message, result_type)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", id=id, error=e.error)
+
+        # The message coming here is the generic request. We look at it's .method
+        # to determine which client function to call and to further subparse the params.
+        # Note(zeiler): unfortunately the pydantic types in mcp/types.py are not consistent.
+        # The JSONRPCRequest are supposed to have an id but the InitializeRequest
+        # does not have it.
+        d = json.loads(msg)
+
+        # If we have an id it's a JSONRPCRequest
+        if not d.get('method', None).startswith("notifications/"):
+            # rpc_message2 = types.JSONRPCRequest.model_validate(rpc_message)
+            # underlying: types.JSONRPCRequest = rpc_message.root
+            client_message = types.ClientRequest.model_validate(d)
+            response = asyncio.run(
+                send_request(client_message, id=d.get('id', ""))
+            )  # underlying.id))
+        else:  # JSONRPCRequest
+            client_message = types.ClientNotification.model_validate(d)
+            response = asyncio.run(send_notification(client_message))
+        if response is None:
+            return ""
+        # return as a serialized json string
+        return response.model_dump_json(by_alias=True, exclude_none=True)

clarifai/runners/models/model_builder.py

@@ -410,11 +410,11 @@ class ModelBuilder:
         signatures = {method.name: method.signature for method in method_info.values()}
         return signatures_to_yaml(signatures)
 
-    def get_method_signatures(self):
+    def get_method_signatures(self, mocking=True):
         """
         Returns the method signatures for the model class.
         """
-        model_class = self.load_model_class(mocking=True)
+        model_class = self.load_model_class(mocking=mocking)
         method_info = model_class._get_method_info()
         signatures = [method.signature for method in method_info.values()]
         return signatures
@@ -934,18 +934,43 @@ class ModelBuilder:
         try:
             is_uploaded = self.monitor_model_build()
             if is_uploaded:
-                from clarifai.runners.utils import code_script
-
-                method_signatures = self.get_method_signatures()
-                snippet = code_script.generate_client_script(
-                    method_signatures,
-                    user_id=self.client.user_app_id.user_id,
-                    app_id=self.client.user_app_id.app_id,
-                    model_id=self.model_proto.id,
-                )
+                # Provide an mcp client config
+                if model_type_id == "mcp":
+                    snippet = (
+                        """
+import asyncio
+import os
+from fastmcp import Client
+from fastmcp.client.transports import StreamableHttpTransport
+
+transport = StreamableHttpTransport(url="%s/mcp",
+                                    headers={"Authorization": "Bearer " + os.environ["CLARIFAI_PAT"]})
+
+async def main():
+  async with Client(transport) as client:
+    tools = await client.list_tools()
+    print(f"Available tools: {tools}")
+    result = await client.call_tool(tools[0].name, {"a": 5, "b": 3})
+    print(f"Result: {result[0].text}")
+
+if __name__ == "__main__":
+  asyncio.run(main())
+"""
+                        % self.model_url
+                    )
+                else:  # python code to run the model.
+                    from clarifai.runners.utils import code_script
+
+                    method_signatures = self.get_method_signatures()
+                    snippet = code_script.generate_client_script(
+                        method_signatures,
+                        user_id=self.client.user_app_id.user_id,
+                        app_id=self.client.user_app_id.app_id,
+                        model_id=self.model_proto.id,
+                    )
                 logger.info("""\n
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
-# Here is a code snippet to call this model:
+# Here is a code snippet to use this model:
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                 """)
                 logger.info(snippet)

clarifai/runners/models/test_model_builder.py

@@ -0,0 +1,89 @@
+import unittest
+from unittest.mock import patch, mock_open
+from clarifai.runners.models.model_builder import (
+    ModelBuilder,
+    DEFAULT_DOWNLOAD_CHECKPOINT_WHEN,
+)
+
+mock_data = """
+checkpoints:
+  type: huggingface
+  repo_id: test_repo
+model:
+  user_id: test_user
+  app_id: test_app
+  id: test_model
+  model_type_id: test_type
+inference_compute_info:
+  cpu_limit: "2"
+  cpu_memory: "4Gi"
+  num_accelerators: 0
+"""
+
+
+class TestModelBuilder(unittest.TestCase):
+    @patch("builtins.open", new_callable=mock_open, read_data=mock_data)
+    @patch("os.path.exists")
+    @patch("os.listdir")
+    def test_validate_folder(self, mock_listdir, mock_exists, mock_open_file):
+        mock_exists.return_value = True
+        mock_listdir.return_value = ["config.yaml", "1", "requirements.txt"]
+
+        builder = ModelBuilder(folder="test_folder", validate_api_ids=False)
+        validated_folder = builder._validate_folder("test_folder")
+        self.assertIn("test_folder", validated_folder)
+
+    def test_default_download_checkpoint_when(self):
+        self.assertEqual(DEFAULT_DOWNLOAD_CHECKPOINT_WHEN, "runtime")
+
+    @patch(
+        "builtins.open",
+        new_callable=mock_open,
+        read_data="checkpoints:\n  type: huggingface\n  repo_id: test_repo\n",
+    )
+    @patch("os.path.exists")
+    @patch("yaml.safe_load")
+    def test_load_config(self, mock_yaml_load, mock_exists, mock_open_file):
+        mock_exists.return_value = True
+        mock_yaml_load.return_value = {
+            "checkpoints": {"type": "huggingface", "repo_id": "test_repo"}
+        }
+        config = ModelBuilder._load_config("config.yaml")
+        self.assertIn("checkpoints", config)
+        self.assertEqual(config["checkpoints"]["type"], "huggingface")
+
+    @patch("shutil.copy")
+    @patch("os.path.exists")
+    def test_backup_config(self, mock_exists, mock_copy):
+        mock_exists.side_effect = [True, False]  # config exists, backup does not
+        ModelBuilder._backup_config("config.yaml")
+        mock_copy.assert_called_once_with("config.yaml", "config.yaml.bak")
+
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("yaml.safe_dump")
+    def test_save_config(self, mock_yaml_dump, mock_open_file):
+        config = {"key": "value"}
+        ModelBuilder._save_config("config.yaml", config)
+        mock_yaml_dump.assert_called_once_with(config, mock_open_file())
+
+    def test_validate_config_model(self):
+        builder = ModelBuilder(folder="test_folder", validate_api_ids=False)
+        builder.config = {"model": {}}
+        try:
+            builder._validate_config_model()
+        except AssertionError:
+            self.fail("_validate_config_model raised AssertionError unexpectedly!")
+
+    @patch("os.path.exists")
+    @patch("os.listdir")
+    def test_validate_folder_missing_config(self, mock_listdir, mock_exists):
+        mock_exists.return_value = True
+        mock_listdir.return_value = ["1", "requirements.txt"]
+
+        builder = ModelBuilder(folder="test_folder", validate_api_ids=False)
+        with self.assertRaises(AssertionError):
+            builder._validate_folder("test_folder")
+
+
+if __name__ == "__main__":
+    unittest.main()