add unit test

Author: [email protected]

mlx_lm/server.py

@@ -406,17 +406,10 @@ def do_POST(self):
             self.body = json.loads(raw_body.decode())
         except json.JSONDecodeError as e:
             logging.error(f"JSONDecodeError: {e} - Raw body: {raw_body.decode()}")
-            # Set appropriate headers based on streaming requirement
-            if self.stream:
-                self._set_stream_headers(400)
-                self.wfile.write(
-                    f"data: {json.dumps({'error': f'Invalid JSON in request body: {e}'})}\n\n".encode()
-                )
-            else:
-                self._set_completion_headers(400)
-                self.wfile.write(
-                    json.dumps({"error": f"Invalid JSON in request body: {e}"}).encode()
-                )
+            self._set_completion_headers(400)
+            self.wfile.write(
+                json.dumps({"error": f"Invalid JSON in request body: {e}"}).encode()
+            )
             return
 
         indent = "\t"  # Backslashes can't be inside of f-strings
@@ -436,8 +429,12 @@ def do_POST(self):
             model_path = os.environ['MLX_MODEL_PATH']
             if not os.path.exists(model_path):
                 raise Exception(f"MLX_MODEL_PATH={model_path} is not a path")
-            self.requested_model = os.path.join(model_path, self.requested_model)
-            self.requested_draft_model = os.path.join(model_path, self.requested_draft_model)
+
+            if self.requested_model != "default_model":
+                self.requested_model = os.path.join(model_path, self.requested_model)
+
+            if self.requested_draft_model != "default_model":
+                self.requested_draft_model = os.path.join(model_path, self.requested_draft_model)
 
         self.num_draft_tokens = self.body.get(
             "num_draft_tokens", self.model_provider.cli_args.num_draft_tokens

tests/test_server.py

@@ -7,6 +7,7 @@
 import unittest
 
 import requests
+import time
 
 from mlx_lm.server import APIHandler
 from mlx_lm.utils import load
@@ -69,6 +70,55 @@ def tearDownClass(cls):
         cls.httpd.server_close()
         cls.server_thread.join()
 
+    def test_handle_chunked_request(self):
+        url = f"http://localhost:{self.port}/v1/chat/completions"
+
+        post_data = {
+            "model": "default_model",
+            "prompt": "Once upon a time",
+            "max_tokens": 10,
+            "temperature": 0.0,
+            "stream": False,
+            "top_p": 1.0,
+        }
+
+        # chunked request
+        data_parts = [
+            b'{"model": "default_model", "messages": [{"role": "user", "content": "Once',
+            b' upon a times, Once upon ',
+            b'a time"}], "temperature": 0.8, "max_tokens": 1024, "stream": false}',
+        ]
+
+        max_length = 0
+        for part in data_parts:
+            max_length += len(part)
+
+        def data_generator():
+            for part in data_parts:
+                yield part
+                time.sleep(0.1)
+
+        try:
+            response = requests.post(
+                url,
+                data=data_generator(),
+                headers={
+                    "Transfer-Encoding": "chunked",
+                    "Content-Type": "application/json",
+                },
+            )
+            self.assertEqual(response.status_code, 200)
+        except requests.exceptions.RequestException:
+            self.assertTrue(False, "Chunked request failed")
+
+        response_body = json.loads(response.text)
+        self.assertIn("id", response_body)
+        self.assertIn("choices", response_body)
+        self.assertIn("usage", response_body) 
+
+        # Check that tokens were generated
+        self.assertTrue(response_body["usage"]["completion_tokens"] > 0)
+
     def test_handle_completions(self):
         url = f"http://localhost:{self.port}/v1/completions"