diff --git a/prospector/llm/models/gemini.py b/prospector/llm/models/gemini.py index 4ab85c652..147086254 100644 --- a/prospector/llm/models/gemini.py +++ b/prospector/llm/models/gemini.py @@ -58,15 +58,29 @@ def _call( ], } - response = requests.post(endpoint, headers=headers, json=data) - - if not response.status_code == 200: + try: + response = requests.post(endpoint, headers=headers, json=data) + return self.parse(response.json()) + except requests.exceptions.HTTPError as http_error: logger.error( - f"Invalid response from AI Core API with error code {response.status_code}" + f"HTTP error occurred when sending a request through AI Core: {http_error}" ) - raise Exception("Invalid response from AI Core API.") - - return self.parse(response.json()) + raise + except requests.exceptions.Timeout as timeout_err: + logger.error( + f"Timeout error occured when sending a request through AI Core: {timeout_err}" + ) + raise + except requests.exceptions.ConnectionError as conn_err: + logger.error( + f"Connection error occurred when sending a request through AI Core: {conn_err}" + ) + raise + except requests.exceptions.RequestException as req_err: + logger.error( + f"A request error occured when sending a request through AI Core: {req_err}" + ) + raise def parse(self, message) -> str: """Parse the returned JSON object from OpenAI.""" diff --git a/prospector/llm/models/mistral.py b/prospector/llm/models/mistral.py index a413fc316..9708d8e31 100644 --- a/prospector/llm/models/mistral.py +++ b/prospector/llm/models/mistral.py @@ -39,15 +39,29 @@ def _call( "messages": [{"role": "user", "content": prompt}], } - response = requests.post(endpoint, headers=headers, json=data) - - if not response.status_code == 200: + try: + response = requests.post(endpoint, headers=headers, json=data) + return self.parse(response.json()) + except requests.exceptions.HTTPError as http_error: logger.error( - f"Invalid response from AI Core API with error code {response.status_code}" + f"HTTP error occurred when sending a request through AI Core: {http_error}" ) - raise Exception("Invalid response from AI Core API.") - - return self.parse(response.json()) + raise + except requests.exceptions.Timeout as timeout_err: + logger.error( + f"Timeout error occured when sending a request through AI Core: {timeout_err}" + ) + raise + except requests.exceptions.ConnectionError as conn_err: + logger.error( + f"Connection error occurred when sending a request through AI Core: {conn_err}" + ) + raise + except requests.exceptions.RequestException as req_err: + logger.error( + f"A request error occured when sending a request through AI Core: {req_err}" + ) + raise def parse(self, message) -> str: """Parse the returned JSON object from OpenAI.""" diff --git a/prospector/llm/models/openai.py b/prospector/llm/models/openai.py index f3e132dfe..e34b647e0 100644 --- a/prospector/llm/models/openai.py +++ b/prospector/llm/models/openai.py @@ -30,7 +30,8 @@ def _identifying_params(self) -> Dict[str, Any]: def _call( self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any ) -> str: - endpoint = f"{self.deployment_url}/chat/completions?api-version=2023-05-15" + # endpoint = f"{self.deployment_url}/chat/completions?api-version=2023-05-15" + endpoint = f"{self.deployment_url}/chat/cpletions?api-version=2023-05-15" headers = instantiation.get_headers(self.ai_core_sk_filepath) data = { "messages": [ @@ -42,15 +43,29 @@ def _call( "temperature": self.temperature, } - response = requests.post(endpoint, headers=headers, json=data) - - if not response.status_code == 200: + try: + response = requests.post(endpoint, headers=headers, json=data) + return self.parse(response.json()) + except requests.exceptions.HTTPError as http_error: logger.error( - f"Invalid response from AI Core API with error code {response.status_code}" + f"HTTP error occurred when sending a request through AI Core: {http_error}" ) - raise Exception("Invalid response from AI Core API.") - - return self.parse(response.json()) + raise + except requests.exceptions.Timeout as timeout_err: + logger.error( + f"Timeout error occured when sending a request through AI Core: {timeout_err}" + ) + raise + except requests.exceptions.ConnectionError as conn_err: + logger.error( + f"Connection error occurred when sending a request through AI Core: {conn_err}" + ) + raise + except requests.exceptions.RequestException as req_err: + logger.error( + f"A request error occured when sending a request through AI Core: {req_err}" + ) + raise def parse(self, message) -> str: """Parse the returned JSON object from OpenAI."""