diff --git a/flowise/client.py b/flowise/client.py
index 5b9fe37..ada6959 100644
--- a/flowise/client.py
+++ b/flowise/client.py
@@ -38,9 +38,10 @@ def __init__(
 
 
 class Flowise:
-    def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None):
+    def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None, timeout: Optional[int] = None):
         self.base_url = base_url or 'http://localhost:3000'
         self.api_key = api_key or ''
+        self.timeout = timeout
 
     def _get_headers(self) -> Dict[str, str]:
         headers = {}
@@ -71,7 +72,7 @@ def create_prediction(self, data: PredictionData) -> Generator[str, None, None]:
                 'uploads': [upload.__dict__ for upload in (data.uploads or [])]
             }
 
-            with requests.post(prediction_url, json=prediction_payload, stream=True, headers=self._get_headers()) as r:
+            with requests.post(prediction_url, json=prediction_payload, stream=True, headers=self._get_headers(), timeout=self.timeout) as r:
                 r.raise_for_status()
                 for line in r.iter_lines():
                     if line:
@@ -91,6 +92,6 @@ def create_prediction(self, data: PredictionData) -> Generator[str, None, None]:
                 'uploads': [upload.__dict__ for upload in (data.uploads or [])]
             }
 
-            response = requests.post(prediction_url, json=prediction_payload, headers=self._get_headers())
+            response = requests.post(prediction_url, json=prediction_payload, headers=self._get_headers(), timeout=self.timeout)
             response.raise_for_status()
             yield response.json()
diff --git a/tests/test_flowise.py b/tests/test_flowise.py
index 6882d32..4f5e985 100644
--- a/tests/test_flowise.py
+++ b/tests/test_flowise.py
@@ -36,10 +36,13 @@ def test_create_prediction_streaming(self, mock_get, mock_post):
         mock_get.return_value.json.return_value = {"isStreaming": True}
 
         # Mock the streaming POST response
-        mock_post.return_value.iter_lines.return_value = [
+        mock_response = MagicMock()
+        mock_response.__enter__.return_value.iter_lines.return_value = [
             b'data: {"event": "token", "data": "Why don\'t scientists trust atoms?"}',
             b'data: {"event": "token", "data": "Because they make up everything!"}'
         ]
+        mock_response.__enter__.return_value.raise_for_status.return_value = None
+        mock_post.return_value = mock_response
 
         # Create a client instance
         client = Flowise()