Missing Request Schema for SageMaker Endpoint in jupyter-ai Extension

[Naveen-Annam-CMA]

Description:

When trying to use a SageMaker Endpoint as the completion model in the jupyter-ai chat UI, we encounter the following error:

Traceback (most recent call last):
  File "/home/ubuntu/miniconda/lib/python3.10/site-packages/jupyter_ai/chat_handlers/base.py", line 196, in on_message
    await self.process_message(message)
  File "/home/ubuntu/miniconda/lib/python3.10/site-packages/jupyter_ai/chat_handlers/default.py", line 96, in process_message
    self.get_llm_chain()
  File "/home/ubuntu/miniconda/lib/python3.10/site-packages/jupyter_ai/chat_handlers/base.py", line 345, in get_llm_chain
    self create_llm_chain(lm_provider, lm_provider_params)
  File "/home/ubuntu/miniconda/lib/python3.10/site-packages/jupyter_ai/chat_handlers/default.py", line 38, in create_llm_chain
    llm = provider(**unified_parameters)
  File "/home/ubuntu/miniconda/lib/python3.10/site-packages/jupyter_ai_magics/partner_providers/aws.py", line 205, in __init__
    request_schema = kwargs.pop("request_schema")
KeyError: 'request_schema'

Error related to a missing request_schema. We have set all the required fields (including the request_schema) but we seem to be encountering the above. It appears that the request schema for the SageMaker Endpoint is missing or not being properly handled by the extension.

Environment:

• Platform: JupyterLab running on AWS with jupyter-ai Extension
• Issue: Missing request_schema when using SageMaker Endpoint as the completion model.

Request for Assistance:

• Please help resolve the missing request_schema issue for SageMaker Endpoint.