Repository: Clarifai/clarifai-python

Issue:
Add an OpenAIModelClass
Add a model class just like was done in this PR https://github.com/Clarifai/clarifai-python/pull/592 for MCPModelClass that uses an openai_transport to transfer the compatible request and response formats for an openAI compatible API client to the model.py that a developer can fill in. They can subclass OpenAI and provide an openAI compatible server that the openai_transport will interact with. Add a dummy model like we did for mcp, and add some test cases for both the MCPModelClass and OpenAIModelClass.

Discussion:

Solution patch:
diff --git a/clarifai/__init__.py b/clarifai/__init__.py
index 24cf2a7c..4bcb58b4 100644
--- a/clarifai/__init__.py
+++ b/clarifai/__init__.py
@@ -1 +1 @@
-__version__ = "11.4.3"
+__version__ = "11.4.4"
diff --git a/clarifai/client/model_client.py b/clarifai/client/model_client.py
index e5ec431b..c576a230 100644
--- a/clarifai/client/model_client.py
+++ b/clarifai/client/model_client.py
@@ -196,7 +196,7 @@ def f(*args, **kwargs):
             f.__name__ = method_name
             f.__qualname__ = f'{self.__class__.__name__}.{method_name}'
             input_annotations = code_script._get_annotations_source(method_signature)
-            return_annotation = input_annotations.pop('return', (None, None))[0]
+            return_annotation = input_annotations.pop('return', (None, None, None))[0]
             sig = inspect.signature(f).replace(
                 parameters=[
                     inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v[0])
diff --git a/clarifai/runners/__init__.py b/clarifai/runners/__init__.py
index 8fa0b8ef..fd60f286 100644
--- a/clarifai/runners/__init__.py
+++ b/clarifai/runners/__init__.py
@@ -1,9 +1,13 @@
+from .models.mcp_class import MCPModelClass
 from .models.model_builder import ModelBuilder
 from .models.model_class import ModelClass
 from .models.model_runner import ModelRunner
+from .models.openai_class import OpenAIModelClass
 
 __all__ = [
     "ModelRunner",
     "ModelBuilder",
     "ModelClass",
+    "MCPModelClass",
+    "OpenAIModelClass",
 ]
diff --git a/clarifai/runners/models/__init__.py b/clarifai/runners/models/__init__.py
index e69de29b..9a98497b 100644
--- a/clarifai/runners/models/__init__.py
+++ b/clarifai/runners/models/__init__.py
@@ -0,0 +1,9 @@
+from clarifai.runners.models.mcp_class import MCPModelClass
+from clarifai.runners.models.model_class import ModelClass
+from clarifai.runners.models.openai_class import OpenAIModelClass
+
+__all__ = [
+    "ModelClass",
+    "MCPModelClass",
+    "OpenAIModelClass",
+]
diff --git a/clarifai/runners/models/dummy_openai_model.py b/clarifai/runners/models/dummy_openai_model.py
new file mode 100644
index 00000000..23d18fb9
--- /dev/null
+++ b/clarifai/runners/models/dummy_openai_model.py
@@ -0,0 +1,148 @@
+"""Dummy OpenAI model implementation for testing."""
+
+import json
+from typing import Iterator
+
+from clarifai.runners.models.openai_class import OpenAIModelClass
+
+
+class MockOpenAIClient:
+    """Mock OpenAI client for testing."""
+
+    class Completions:
+        def create(self, **kwargs):
+            """Mock create method for compatibility."""
+            if kwargs.get("stream", False):
+                return MockCompletionStream(kwargs.get("messages", []))
+            else:
+                return MockCompletion(kwargs.get("messages", []))
+
+    def __init__(self):
+        self.chat = self  # Make self.chat point to self for compatibility
+        self.completions = self.Completions()  # For compatibility with some clients
+
+
+class MockCompletion:
+    """Mock completion object that mimics the OpenAI completion response structure."""
+
+    class Choice:
+        class Message:
+            def __init__(self, content):
+                self.content = content
+                self.role = "assistant"
+
+        def __init__(self, content):
+            self.message = self.Message(content)
+            self.finish_reason = "stop"
+            self.index = 0
+
+    def __init__(self, messages):
+        # Generate a simple response based on the last message
+        last_message = messages[-1] if messages else {"content": ""}
+        response_text = f"Echo: {last_message.get('content', '')}"
+
+        self.choices = [self.Choice(response_text)]
+        self.usage = {
+            "prompt_tokens": len(str(messages)),
+            "completion_tokens": len(response_text),
+            "total_tokens": len(str(messages)) + len(response_text),
+        }
+        self.id = "dummy-completion-id"
+        self.created = 1234567890
+        self.model = "dummy-model"
+
+
+class MockCompletionStream:
+    """Mock completion stream that mimics the OpenAI streaming response structure."""
+
+    class Chunk:
+        class Choice:
+            class Delta:
+                def __init__(self, content=None):
+                    self.content = content
+                    self.role = "assistant" if content is None else None
+
+            def __init__(self, content=None):
+                self.delta = self.Delta(content)
+                self.finish_reason = None if content else "stop"
+                self.index = 0
+
+        def __init__(self, content=None):
+            self.choices = [self.Choice(content)]
+            self.id = "dummy-chunk-id"
+            self.created = 1234567890
+            self.model = "dummy-model"
+
+    def __init__(self, messages):
+        # Generate a simple response based on the last message
+        last_message = messages[-1] if messages else {"content": ""}
+        self.response_text = f"Echo: {last_message.get('content', '')}"
+        # Divide the response into chunks of 5 characters
+        self.chunks = [self.response_text[i : i + 5] for i in range(0, len(self.response_text), 5)]
+        self.current_chunk = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.current_chunk < len(self.chunks):
+            chunk = self.Chunk(self.chunks[self.current_chunk])
+            self.current_chunk += 1
+            return chunk
+        elif self.current_chunk == len(self.chunks):
+            # Final chunk with empty content to indicate completion
+            self.current_chunk += 1
+            return self.Chunk()
+        else:
+            raise StopIteration
+
+
+class DummyOpenAIModel(OpenAIModelClass):
+    """Dummy OpenAI model implementation for testing."""
+
+    def get_openai_client(self):
+        """Return a mock OpenAI client."""
+        return MockOpenAIClient()
+
+    def _process_request(self, model, messages, temperature=1.0, max_tokens=None):
+        """Process a request for non-streaming responses."""
+        # Simply return the text of the last message as our "response"
+        last_message = messages[-1] if messages else {"content": ""}
+        return f"Echo: {last_message.get('content', '')}"
+
+    def _process_streaming_request(self, model, messages, temperature=1.0, max_tokens=None):
+        """Process a request for streaming responses."""
+        # Generate a simple text response
+        last_message = messages[-1] if messages else {"content": ""}
+        response_text = f"Echo: {last_message.get('content', '')}"
+
+        # Yield chunks of the response
+        for i in range(0, len(response_text), 5):
+            yield response_text[i : i + 5]
+
+    # Override the method directly for testing
+    @OpenAIModelClass.method
+    def openai_stream_transport(self, req: str) -> Iterator[str]:
+        """Direct implementation for testing purposes."""
+        try:
+            request_data = json.loads(req)
+            messages = request_data.get("messages", [])
+
+            # Validate the request format for testing
+            if not messages or not isinstance(messages[0], dict) or "content" not in messages[0]:
+                raise ValueError("Invalid message format")
+
+            last_message = messages[-1] if messages else {"content": ""}
+            response_text = f"Echo: {last_message.get('content', '')}"
+
+            chunks = [response_text[i : i + 5] for i in range(0, len(response_text), 5)]
+            for chunk in chunks:
+                yield chunk
+        except Exception as e:
+            yield f"Error: {str(e)}"
+
+    # Additional example method that could be added for specific model implementations
+    @OpenAIModelClass.method
+    def test_method(self, prompt: str) -> str:
+        """Test method that simply echoes the input."""
+        return f"Test: {prompt}"
diff --git a/clarifai/runners/models/openai_class.py b/clarifai/runners/models/openai_class.py
new file mode 100644
index 00000000..452a720a
--- /dev/null
+++ b/clarifai/runners/models/openai_class.py
@@ -0,0 +1,147 @@
+"""Base class for creating OpenAI-compatible API server."""
+
+import json
+from typing import Any, Dict, Iterator, List, Optional
+
+from clarifai.runners.models.model_class import ModelClass
+from clarifai.runners.utils.openai_convertor import openai_response
+
+
+class OpenAIModelClass(ModelClass):
+    """Base class for wrapping OpenAI-compatible servers as a model running in Clarifai.
+    This handles all the transport between the API and the OpenAI-compatible server.
+    Simply subclass this and implement the get_openai_client() method to return
+    the OpenAI-compatible client instance. The client is then used to handle all
+    the requests and responses.
+    """
+
+    def load_model(self):
+        """Initialize the OpenAI client."""
+        self.client = self.get_openai_client()
+
+    def get_openai_client(self) -> Any:
+        """Required method for each subclass to implement to return the OpenAI-compatible client to use."""
+        raise NotImplementedError("Subclasses must implement get_openai_client() method")
+
+    @ModelClass.method
+    def openai_transport(self, msg: str) -> str:
+        """The single model method to get the OpenAI-compatible request and send it to the OpenAI server
+        then return its response.
+        """
+        # Parse the incoming message as JSON
+        request_data = json.loads(msg)
+
+        # Extract key parameters from the request
+        model = request_data.get("model", "")
+        messages = request_data.get("messages", [])
+        stream = request_data.get("stream", False)
+        temperature = request_data.get("temperature", 1.0)
+        max_tokens = request_data.get("max_tokens", None)
+
+        # Process the request using the provided OpenAI client
+        try:
+            if stream:
+                # For streaming responses
+                chunks = self._process_streaming_request(
+                    model=model, messages=messages, temperature=temperature, max_tokens=max_tokens
+                )
+                # Format the response as an OpenAI streaming response
+                response = openai_response(generated_text=chunks, model=model, stream=True)
+                # Since this is a generator, we need to convert it to a list for JSON serialization
+                response_list = list(response)
+                return json.dumps(response_list)
+            else:
+                # For non-streaming responses
+                completion = self._process_request(
+                    model=model, messages=messages, temperature=temperature, max_tokens=max_tokens
+                )
+
+                # Format the response as an OpenAI response
+                response = openai_response(generated_text=completion, model=model, stream=False)
+                return json.dumps(response)
+        except Exception as e:
+            error_response = {
+                "error": {
+                    "message": str(e),
+                    "type": "InvalidRequestError",
+                    "code": "invalid_request_error",
+                }
+            }
+            return json.dumps(error_response)
+
+    @ModelClass.method
+    def openai_stream_transport(self, req: str) -> Iterator[str]:
+        """Process an OpenAI-compatible request and return a streaming response iterator.
+
+        This method is used when stream=True and returns an iterator of strings directly,
+        without converting to a list or JSON serializing.
+
+        Args:
+            req: The request as a JSON string.
+
+        Returns:
+            Iterator[str]: An iterator yielding text chunks from the streaming response.
+        """
+        # Parse the incoming message
+        request_data = json.loads(req)
+
+        # Extract key parameters from the request
+        model = request_data.get("model", "")
+        messages = request_data.get("messages", [])
+        temperature = request_data.get("temperature", 1.0)
+        max_tokens = request_data.get("max_tokens", None)
+
+        # Process the streaming request and return the iterator directly
+        try:
+            return self._process_streaming_request(
+                model=model, messages=messages, temperature=temperature, max_tokens=max_tokens
+            )
+        except Exception as e:
+            # For errors, yield a single error message
+            yield f"Error: {str(e)}"
+
+    def _process_request(
+        self,
+        model: str,
+        messages: List[Dict[str, Any]],
+        temperature: float = 1.0,
+        max_tokens: Optional[int] = None,
+    ) -> str:
+        """Process a standard (non-streaming) request using the OpenAI client.
+
+        Override this method in your subclass if you need custom processing logic.
+        """
+        # Default implementation - subclasses should override with actual client interaction
+        completion_args = {"model": model, "messages": messages, "temperature": temperature}
+        if max_tokens is not None:
+            completion_args["max_tokens"] = max_tokens
+
+        completion = self.client.chat.completions.create(**completion_args)
+        return completion.choices[0].message.content
+
+    def _process_streaming_request(
+        self,
+        model: str,
+        messages: List[Dict[str, Any]],
+        temperature: float = 1.0,
+        max_tokens: Optional[int] = None,
+    ) -> Iterator[str]:
+        """Process a streaming request using the OpenAI client.
+
+        Override this method in your subclass if you need custom streaming logic.
+        """
+        # Default implementation - subclasses should override with actual client interaction
+        completion_args = {
+            "model": model,
+            "messages": messages,
+            "temperature": temperature,
+            "stream": True,
+        }
+        if max_tokens is not None:
+            completion_args["max_tokens"] = max_tokens
+
+        completion_stream = self.client.chat.completions.create(**completion_args)
+        for chunk in completion_stream:
+            if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
+                if chunk.choices[0].delta.content:
+                    yield chunk.choices[0].delta.content
diff --git a/clarifai/runners/utils/code_script.py b/clarifai/runners/utils/code_script.py
index 3d26946f..1f2a3511 100644
--- a/clarifai/runners/utils/code_script.py
+++ b/clarifai/runners/utils/code_script.py
@@ -87,13 +87,13 @@ async def main():
         method_name = method_signature.name
         client_script_str = f'response = model.{method_name}('
         annotations = _get_annotations_source(method_signature)
-        for param_name, (param_type, default_value) in annotations.items():
+        for param_name, (param_type, default_value, required) in annotations.items():
             print(
                 f"param_name: {param_name}, param_type: {param_type}, default_value: {default_value}"
             )
             if param_name == "return":
                 continue
-            if default_value is None:
+            if default_value is None and required:
                 default_value = _set_default_value(param_type)
             client_script_str += f"{param_name}={default_value}, "
         client_script_str = client_script_str.rstrip(", ") + ")"
@@ -132,7 +132,7 @@ def _get_annotations_source(method_signature: resources_pb2.MethodSignature) ->
         default_value = None
         if data_utils.Param.get_default(input_field):
             default_value = _parse_default_value(input_field)
-        annotations[param_name] = (param_type, default_value)
+        annotations[param_name] = (param_type, default_value, input_field.required)
     if not method_signature.output_fields:
         raise ValueError("MethodSignature must have at least one output field")
     for output_field in method_signature.output_fields:
@@ -140,7 +140,7 @@ def _get_annotations_source(method_signature: resources_pb2.MethodSignature) ->
         param_type = _get_base_type(output_field)
         if output_field.iterator:
             param_type = f"Iterator[{param_type}]"
-        annotations[param_name] = (param_type, None)
+        annotations[param_name] = (param_type, None, output_field.required)
     return annotations
 
 
@@ -189,7 +189,7 @@ def _map_default_value(field_type):
     default_value = None
 
     if field_type == "str":
-        default_value = 'What is the future of AI?'
+        default_value = repr('What is the future of AI?')
     elif field_type == "bytes":
         default_value = b""
     elif field_type == "int":
@@ -224,11 +224,9 @@ def _set_default_value(field_type):
     Set the default value of a field if it is not set.
     """
     is_iterator = False
-    print(f"before field_type: {field_type}")
     if field_type.startswith("Iterator["):
         is_iterator = True
         field_type = field_type[9:-1]
-    print(f"after field_type: {field_type}")
     default_value = None
     default_value = _map_default_value(field_type)
     if field_type.startswith("List["):
@@ -245,11 +243,8 @@ def _set_default_value(field_type):
         element_type_defaults = [_map_default_value(et) for et in element_types]
         default_value = f"{{{', '.join([str(et) for et in element_type_defaults])}}}"
 
-    if field_type == 'str':
-        default_value = repr(default_value)
     if is_iterator:
         default_value = f'iter([{default_value}])'
-    print(f"after default_value: {default_value}")
     return default_value
 
 
diff --git a/clarifai/runners/utils/data_utils.py b/clarifai/runners/utils/data_utils.py
index db6003f0..c8725c7b 100644
--- a/clarifai/runners/utils/data_utils.py
+++ b/clarifai/runners/utils/data_utils.py
@@ -320,9 +320,7 @@ def to_proto(self, proto=None) -> ParamProto:
             proto.model_type_range_info.CopyFrom(range_info)
         proto.is_param = self.is_param
 
-        if self.default is not None:
-            proto = self.set_default(proto, self.default)
-
+        proto = self.set_default(proto, self.default)
         return proto
 
     @classmethod
diff --git a/clarifai/runners/utils/method_signatures.py b/clarifai/runners/utils/method_signatures.py
index 36ca8680..4b452b95 100644
--- a/clarifai/runners/utils/method_signatures.py
+++ b/clarifai/runners/utils/method_signatures.py
@@ -74,32 +74,12 @@ def build_function_signature(func):
     assert method_type in ('UNARY_UNARY', 'UNARY_STREAMING', 'STREAMING_STREAMING')
     # method_signature.method_type = method_type
     method_signature.description = inspect.cleandoc(func.__doc__ or '')
-    # method_signature.annotations_json = json.dumps(_get_annotations_source(func))
 
     method_signature.input_fields.extend(input_sigs)
     method_signature.output_fields.append(output_sig)
     return method_signature
 
 
-# def _get_annotations_source(func):
-#   """Extracts raw annotation strings from the function source."""
-#   source = inspect.getsource(func)  # Get function source code
-#   source = textwrap.dedent(source)  # Dedent source code
-#   tree = ast.parse(source)  # Parse into AST
-#   func_node = next(node for node in tree.body
-#                    if isinstance(node, ast.FunctionDef))  # Get function node
-
-#   annotations = {}
-#   for arg in func_node.args.args:  # Process arguments
-#     if arg.annotation:
-#       annotations[arg.arg] = ast.unparse(arg.annotation)  # Get raw annotation string
-
-#   if func_node.returns:  # Process return type
-#     annotations["return"] = ast.unparse(func_node.returns)
-
-#   return annotations
-
-
 def _process_input_field(field: resources_pb2.ModelTypeField) -> str:
     base_type = _get_base_type(field)
     if field.iterator:
diff --git a/clarifai/runners/utils/openai_convertor.py b/clarifai/runners/utils/openai_convertor.py
index a194be88..3d682b32 100644
--- a/clarifai/runners/utils/openai_convertor.py
+++ b/clarifai/runners/utils/openai_convertor.py
@@ -171,14 +171,14 @@ def openai_to_hf_messages(openai_messages):
 
 
 def build_openai_messages(
-    prompt: str,
-    image: Image,
-    images: List[Image],
-    audio: Audio,
-    audios: List[Audio],
-    video: Video,
-    videos: List[Video],
-    messages: List[Dict],
+    prompt: str = None,
+    image: Image = None,
+    images: List[Image] = None,
+    audio: Audio = None,
+    audios: List[Audio] = None,
+    video: Video = None,
+    videos: List[Video] = None,
+    messages: List[Dict] = None,
 ) -> List[Dict]:
     """
     Construct OpenAI-compatible messages from input components.


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..e9da29fa
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,21 @@
+"""Conftest for OpenAI tests."""
+
+import sys
+from unittest import mock
+
+# Create mock modules
+mock_fastmcp = mock.MagicMock()
+mock_fastmcp.Client = mock.MagicMock()
+mock_fastmcp.FastMCP = mock.MagicMock()
+
+mock_mcp = mock.MagicMock()
+mock_mcp.types = mock.MagicMock()
+mock_mcp.shared = mock.MagicMock()
+mock_mcp.shared.exceptions = mock.MagicMock()
+mock_mcp.shared.exceptions.McpError = Exception
+
+# Mock the fastmcp and mcp modules
+sys.modules['fastmcp'] = mock_fastmcp
+sys.modules['mcp'] = mock_mcp
+sys.modules['mcp.shared'] = mock_mcp.shared
+sys.modules['mcp.shared.exceptions'] = mock_mcp.shared.exceptions
diff --git a/tests/openai_model_test.py b/tests/openai_model_test.py
new file mode 100644
index 00000000..9d2e81c8
--- /dev/null
+++ b/tests/openai_model_test.py
@@ -0,0 +1,83 @@
+"""Tests for OpenAI class.
+
+This test uses conftest.py in the same directory to set up mocks for dependencies.
+"""
+
+import json
+import os
+import sys
+
+# Import pytest after sys.modules updates in conftest.py
+import pytest
+
+# Add the base directory to the path to allow direct imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))
+
+# Now we can import our module
+from clarifai.runners.models.dummy_openai_model import DummyOpenAIModel
+from clarifai.runners.models.model_class import ModelClass
+from clarifai.runners.models.openai_class import OpenAIModelClass
+
+
+class TestOpenAIModelClass:
+    def test_inheritance(self):
+        """Test that OpenAIModelClass inherits from ModelClass."""
+        assert issubclass(OpenAIModelClass, ModelClass)
+
+    def test_abstract_method(self):
+        """Test that get_openai_client is abstract."""
+        with pytest.raises(NotImplementedError):
+            OpenAIModelClass().get_openai_client()
+
+    def test_dummy_model(self):
+        """Test that DummyOpenAIModel works."""
+        model = DummyOpenAIModel()
+        assert isinstance(model, OpenAIModelClass)
+
+        client = model.get_openai_client()
+        assert client is not None
+
+    def test_transport_method_non_streaming(self):
+        """Test the openai_transport method with non-streaming."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        request = {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "Hello world"}],
+            "stream": False,
+        }
+
+        response = model.openai_transport(json.dumps(request))
+        data = json.loads(response)
+
+        assert "choices" in data
+        assert len(data["choices"]) > 0
+        assert "message" in data["choices"][0]
+        assert "content" in data["choices"][0]["message"]
+
+    def test_transport_method_streaming(self):
+        """Test the openai_transport method with streaming."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        request = {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "Hello world"}],
+            "stream": True,
+        }
+
+        response = model.openai_transport(json.dumps(request))
+        data = json.loads(response)
+
+        assert isinstance(data, list)
+        assert len(data) > 0
+
+        # Check format of streaming chunks
+        for chunk in data[:-1]:  # All except last
+            assert "choices" in chunk
+            assert "delta" in chunk["choices"][0]
+
+        # Check last chunk has finish_reason
+        assert "choices" in data[-1]
+        assert data[-1]["choices"][0]["finish_reason"] == "stop"
diff --git a/tests/requirements.txt b/tests/requirements.txt
index a614933b..79701024 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -7,3 +7,4 @@ pypdf==3.17.4
 seaborn==0.13.2
 pycocotools>=2.0.7
 rich>=13.4.2
+fastmcp>=2.3.4
diff --git a/tests/runners/dummy_runner_models/requirements.txt b/tests/runners/dummy_runner_models/requirements.txt
index f7f3193c..91646b8a 100644
--- a/tests/runners/dummy_runner_models/requirements.txt
+++ b/tests/runners/dummy_runner_models/requirements.txt
@@ -2,3 +2,4 @@ aiohttp
 requests
 clarifai
 huggingface_hub
+fastmcp>=2.3.4
diff --git a/tests/runners/hf_mbart_model/requirements.txt b/tests/runners/hf_mbart_model/requirements.txt
index d9a4d1c1..5ca09138 100644
--- a/tests/runners/hf_mbart_model/requirements.txt
+++ b/tests/runners/hf_mbart_model/requirements.txt
@@ -7,3 +7,4 @@ tiktoken==0.9.0
 tokenizers==0.21.1
 torch==2.6.0
 transformers==4.51.3
+fastmcp>=2.3.4
diff --git a/tests/runners/test_model_classes.py b/tests/runners/test_model_classes.py
new file mode 100644
index 00000000..94432e00
--- /dev/null
+++ b/tests/runners/test_model_classes.py
@@ -0,0 +1,103 @@
+"""Test cases for MCPModelClass and OpenAIModelClass."""
+
+import json
+
+import pytest
+
+from clarifai.runners.models.dummy_openai_model import DummyOpenAIModel
+from clarifai.runners.models.mcp_class import MCPModelClass
+from clarifai.runners.models.openai_class import OpenAIModelClass
+
+
+class TestModelClasses:
+    """Tests for model classes."""
+
+    def test_mcp_model_initialization(self):
+        """Test that MCPModelClass requires subclass implementation."""
+        # Test that subclass must implement get_server()
+        with pytest.raises(NotImplementedError):
+            MCPModelClass().get_server()
+
+    def test_openai_model_initialization(self):
+        """Test that OpenAIModelClass can be initialized."""
+        model = DummyOpenAIModel()
+        assert isinstance(model, OpenAIModelClass)
+
+        # Test that subclass must implement get_openai_client()
+        with pytest.raises(NotImplementedError):
+            OpenAIModelClass().get_openai_client()
+
+    def test_openai_transport_non_streaming(self):
+        """Test OpenAI transport method with non-streaming request."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        # Create a simple chat request
+        request = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, world!"},
+            ],
+            "stream": False,
+        }
+
+        # Call the transport method
+        response_str = model.openai_transport(json.dumps(request))
+        response = json.loads(response_str)
+
+        # Verify the response format
+        assert "id" in response
+        assert "object" in response
+        assert "choices" in response
+        assert len(response["choices"]) > 0
+        assert "message" in response["choices"][0]
+        assert "content" in response["choices"][0]["message"]
+        assert "Echo: Hello, world!" in response["choices"][0]["message"]["content"]
+
+    def test_openai_transport_streaming(self):
+        """Test OpenAI transport method with streaming request."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        # Create a simple chat request with streaming
+        request = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, world!"},
+            ],
+            "stream": True,
+        }
+
+        # Call the transport method
+        response_str = model.openai_transport(json.dumps(request))
+        response_chunks = json.loads(response_str)
+
+        # Verify the response format for streaming
+        assert isinstance(response_chunks, list)
+        assert len(response_chunks) > 0
+
+        # Check all chunks except the last one
+        for chunk in response_chunks[:-1]:
+            assert "id" in chunk
+            assert "object" in chunk
+            assert chunk["object"] == "chat.completion.chunk"
+            assert "choices" in chunk
+            assert len(chunk["choices"]) > 0
+            assert "delta" in chunk["choices"][0]
+            assert "content" in chunk["choices"][0]["delta"]
+
+        # Check the last chunk (should have empty delta and finish_reason)
+        last_chunk = response_chunks[-1]
+        assert "choices" in last_chunk
+        assert len(last_chunk["choices"]) > 0
+        assert "delta" in last_chunk["choices"][0]
+        assert len(last_chunk["choices"][0]["delta"]) == 0
+        assert last_chunk["choices"][0]["finish_reason"] == "stop"
+
+    def test_custom_method(self):
+        """Test custom method on the DummyOpenAIModel."""
+        model = DummyOpenAIModel()
+        result = model.test_method("test input")
+        assert result == "Test: test input"
diff --git a/tests/runners/test_openai_model.py b/tests/runners/test_openai_model.py
new file mode 100644
index 00000000..a7ace09a
--- /dev/null
+++ b/tests/runners/test_openai_model.py
@@ -0,0 +1,127 @@
+"""Test for OpenAIModelClass functionality."""
+
+import json
+
+import pytest
+
+# Import directly from the file we created, not through the import chain
+from clarifai.runners.models.dummy_openai_model import DummyOpenAIModel
+from clarifai.runners.models.openai_class import OpenAIModelClass
+
+
+class TestOpenAIModelClass:
+    """Tests for OpenAIModelClass."""
+
+    def test_openai_model_initialization(self):
+        """Test that OpenAIModelClass can be initialized."""
+        model = DummyOpenAIModel()
+        assert isinstance(model, OpenAIModelClass)
+
+        # Test that subclass must implement get_openai_client()
+        with pytest.raises(NotImplementedError):
+            OpenAIModelClass().get_openai_client()
+
+    def test_openai_transport_non_streaming(self):
+        """Test OpenAI transport method with non-streaming request."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        # Create a simple chat request
+        request = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, world!"},
+            ],
+            "stream": False,
+        }
+
+        # Call the transport method
+        response_str = model.openai_transport(json.dumps(request))
+        response = json.loads(response_str)
+
+        # Verify the response format
+        assert "id" in response
+        assert "object" in response
+        assert "choices" in response
+        assert len(response["choices"]) > 0
+        assert "message" in response["choices"][0]
+        assert "content" in response["choices"][0]["message"]
+        assert "Echo: Hello, world!" in response["choices"][0]["message"]["content"]
+
+    def test_openai_transport_streaming(self):
+        """Test OpenAI transport method with streaming request."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        # Create a simple chat request with streaming
+        request = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, world!"},
+            ],
+            "stream": True,
+        }
+
+        # Call the transport method
+        response_str = model.openai_transport(json.dumps(request))
+        response_chunks = json.loads(response_str)
+
+        # Verify the response format for streaming
+        assert isinstance(response_chunks, list)
+        assert len(response_chunks) > 0
+
+        # Check all chunks except the last one
+        for chunk in response_chunks[:-1]:
+            assert "id" in chunk
+            assert "object" in chunk
+            assert chunk["object"] == "chat.completion.chunk"
+            assert "choices" in chunk
+            assert len(chunk["choices"]) > 0
+            assert "delta" in chunk["choices"][0]
+            assert "content" in chunk["choices"][0]["delta"]
+
+        # Check the last chunk (should have empty delta and finish_reason)
+        last_chunk = response_chunks[-1]
+        assert "choices" in last_chunk
+        assert len(last_chunk["choices"]) > 0
+        assert "delta" in last_chunk["choices"][0]
+        assert len(last_chunk["choices"][0]["delta"]) == 0
+        assert last_chunk["choices"][0]["finish_reason"] == "stop"
+
+    def test_openai_stream_transport(self):
+        """Test the new openai_stream_transport method."""
+        model = DummyOpenAIModel()
+        model.load_model()
+
+        # Create a simple chat request
+        request = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello, world!"},
+            ],
+        }
+
+        # Test with string input
+        response_iter = model.openai_stream_transport(json.dumps(request))
+        chunks = list(response_iter)
+
+        # Verify response format - should be raw text chunks
+        assert len(chunks) > 0
+        combined = ''.join(chunks)
+        assert combined == "Echo: Hello, world!"
+
+        # Test error handling
+        bad_request = json.dumps({"messages": [{"role": "invalid"}]})
+        response_iter = model.openai_stream_transport(bad_request)
+        chunks = list(response_iter)
+        assert len(chunks) == 1
+        assert chunks[0].startswith("Error:")
+
+    def test_custom_method(self):
+        """Test custom method on the DummyOpenAIModel."""
+        model = DummyOpenAIModel()
+        result = model.test_method("test input")
+        assert result == "Test: test input"