jupyterlab / jupyter-ai

A generative AI extension for JupyterLab
https://jupyter-ai.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
3.12k stars 306 forks source link

message.selection in process_message method is None even though a cell is selected. #969

Open anmolbhatia05 opened 2 weeks ago

anmolbhatia05 commented 2 weeks ago

Description

I am trying to create a custom slash command (/observe). The idea is to observe the cell and its output.

in this screenshot, you see that selection is None, even though the cell is selected. Screenshot 2024-08-28 at 18 17 11

My code for the ObserveSlashCommand is as follows:

class ObserveSlashCommand(BaseChatHandler):
    """
    A test slash command implementation that developers should build from. The
    string used to invoke this command is set by the `slash_id` keyword argument
    in the `routing_type` attribute. The command is mainly implemented in the
    `process_message()` method. See built-in implementations under
    `jupyter_ai/handlers` for further reference.

    The provider is made available to Jupyter AI by the entry point declared in
    `pyproject.toml`. If this class or parent module is renamed, make sure the
    update the entry point there as well.
    """

    id = "observe"
    name = "Observe"
    help = "A command to use to observe a particular cell input and output."
    routing_type = SlashCommandRoutingType(slash_id="observe")

    uses_llm = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def create_llm_chain(
        self, provider: Type[BaseProvider], provider_params: Dict[str, str]
    ):
        unified_parameters = {
            **provider_params,
            **(self.get_model_parameters(provider, provider_params)),
        }
        llm = provider(**unified_parameters)

        self.llm = llm
        self.llm_chain = LLMChain(llm=llm, prompt=OBSERVE_PROMPT_TEMPLATE, verbose=True)

    async def process_message(self, message: HumanChatMessage):
        print(message.dict())
        print(message.selection.dict())
        # if not (message.selection and message.selection.type == "cell"):
        #     self.reply(
        #         "`/observe` requires an active code cell without error output. Please click on a cell without error output and retry.",
        #         message,
        #     )
        #     return

        # hint type of selection
        selection: CellSelection = message.selection

        # parse additional instructions specified after `/observe`
        extra_instructions = message.prompt[7:].strip() or "None."

        self.get_llm_chain()
        with self.pending("Analyzing"):
            response = await self.llm_chain.apredict(
                extra_instructions=extra_instructions,
                stop=["\nHuman:"],
                cell_content=selection.source,
                output="temp"
            )
        self.reply(response, message)

Is this enough information ?

Versions -

jupyter_ai                               2.20.0
jupyter_ai_magics                        2.20.0
jupyter_client                           7.4.9
jupyter_core                             5.7.2
jupyter-events                           0.10.0
jupyter-lsp                              2.2.5
jupyter_server                           2.14.2
jupyter_server_terminals                 0.5.3
jupyterlab                               4.2.4
jupyterlab_pygments                      0.3.0
jupyterlab_server                        2.27.3
jupyterlab_widgets                       3.0.13

Reproduce

  1. Go to '...'
  2. Click on '...'
  3. Scroll down to '...'
  4. See error '...'

Expected behavior

Context

Troubleshoot Output
Paste the output from running `jupyter troubleshoot` from the command line here.
You may want to sanitize the paths in the output.
Command Line Output
Paste the output from your command line running `jupyter lab` here, use `--debug` if possible.
Browser Output
Paste the output from your browser Javascript console here, if applicable.

srdas commented 2 weeks ago

@anmolbhatia05 It looks like you took fix.py and made changes. Need to look further but you might want to check also that since your example is a cell with an error you are getting the expected message.selection for cell-with-error. See if your new code works with cells without errors. And a small nit: you may want to change line `extra_instructions = message.prompt[7:].strip() or "None." to integer 8 instead of 7?

anmolbhatia05 commented 2 weeks ago

@srdas Spot on, i took fix.py and tried adapting it to the behaviour i want. I want the cell input output for now and then I can process it the way i want.

I tried it with other cells(that don't have error) but same issue...

I will fix the indexing mistake! Thanks

srdas commented 2 weeks ago

@anmolbhatia05 - Check that it is returning message.selection in /fix in your setup just to make sure you and then debug. Look also at handlers.ts and models.py for reference. Also explain the motivation for /observe in this Issue so that when you get it working, the PR you open has full explanations, to see if this is something other users will use. And finally, add documentation for this in the same PR (you will need to edit /users/index.md and test it with make html).

anmolbhatia05 commented 2 weeks ago

thanks for the suggestions @srdas

print(message.dict())
print(message.selection.dict())

these lines gives the following output -

output from /fix usage -

{'type': 'human', 'id': '839497ff-a429-4a05-8281-a56350363171', 'time': 1724928483.3192031, 'body': '/fix\n\n```\ns\n```\n', 'prompt': '/fix', 'selection': {'type': 'cell-with-error', 'source': 's', 'error': {'name': 'NameError', 'value': "name 's' is not defined", 'traceback': ['\x1b[0;31m---------------------------------------------------------------------------\x1b[0m', '\x1b[0;31mNameError\x1b[0m                                 Traceback (most recent call last)', 'Cell \x1b[0;32mIn[5], line 1\x1b[0m\n\x1b[0;32m----> 1\x1b[0m \x1b[43ms\x1b[49m\n', "\x1b[0;31mNameError\x1b[0m: name 's' is not defined"]}}, 'client': {'username': 'anmol-2', 'initials': 'A', 'name': 'anmol-2', 'display_name': 'anmol-2', 'color': None, 'avatar_url': None, 'id': '8a919c7df7064e0097b23ec755f7d48a'}}
{'type': 'cell-with-error', 'source': 's', 'error': {'name': 'NameError', 'value': "name 's' is not defined", 'traceback': ['\x1b[0;31m---------------------------------------------------------------------------\x1b[0m', '\x1b[0;31mNameError\x1b[0m                                 Traceback (most recent call last)', 'Cell \x1b[0;32mIn[5], line 1\x1b[0m\n\x1b[0;32m----> 1\x1b[0m \x1b[43ms\x1b[49m\n', "\x1b[0;31mNameError\x1b[0m: name 's' is not defined"]}}

output from my /observe usage -

{'type': 'human', 'id': 'a780f920-247a-481d-a827-7e688f533e84', 'time': 1724928008.261506, 'body': '/observe ', 'prompt': '/observe ', 'selection': None, 'client': {'username': 'anmol-2', 'initials': 'A', 'name': 'anmol-2', 'display_name': 'anmol-2', 'color': None, 'avatar_url': None, 'id': '4e8a4861c98e41c7b931b87d30b318ae'}}
[E 2024-08-29 12:40:08.265 AiExtension] 'NoneType' object has no attribute 'dict'

you see that selection is None in observe whereas its populated in fix, can it be that for custom slash commands we have to change something in some other jupyter package ? any hints on how to proceed from here ?

srdas commented 2 weeks ago

@anmolbhatia05 Take a look at how /fix is handled in function onSend in chat-input.tsx where selection is explicitly filled in and you may want to add in logic here to handle /observe in a similar manner. https://github.com/jupyterlab/jupyter-ai/blob/29e2a479e928416e490c4d84b48fac3f421e112a/packages/jupyter-ai/src/components/chat-input.tsx#L187