diff --git a/prompting/llms/apis/sn19_wrapper.py b/prompting/llms/apis/sn19_wrapper.py index a227d2ab5..e8445146b 100644 --- a/prompting/llms/apis/sn19_wrapper.py +++ b/prompting/llms/apis/sn19_wrapper.py @@ -39,6 +39,9 @@ def chat_complete( response = requests.post(url, headers=headers, data=json.dumps(data)) try: response_json = response.json() - return response_json["choices"][0]["message"].get("content") + try: + return response_json["choices"][0]["message"].get("content") + except KeyError: + return response_json["choices"][0]["delta"].get("content") except Exception as e: logger.exception(f"Error in chat_complete: {e}")