diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 0aacd6f6..47e0ca2e 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -27,14 +27,20 @@ class OpenAIBase(): return self._get_response(messages) def _stream_response(self, messages): - response = self.client.chat.completions.create( - model=self.model_name, - messages=messages, - stream=True, - ) - for chunk in response: - if len(chunk.choices) > 0: - yield chunk.choices[0].delta + try: + response = self.client.chat.completions.create( + model=self.model_name, + messages=messages, + stream=True, + ) + for chunk in response: + if len(chunk.choices) > 0: + yield chunk.choices[0].delta + + except Exception as e: + err = f"Error streaming response: {e}, URL: {self.base_url}, API Key: {self.api_key[:5]}***, Model: {self.model_name}" + logger.error(err) + raise Exception(err) def _get_response(self, messages): response = self.client.chat.completions.create(