jupyter-ai icon indicating copy to clipboard operation
jupyter-ai copied to clipboard

Make chat memory size traitlet configurable + /clear to reset memory

Open michaelchia opened this issue 1 year ago • 3 comments
trafficstars

Description

Makes chat memory size configurable via default_max_chat_history traitlet. Also makes /clear reset memory.

The reason why these unrelated features are in same PR stems from the same core change of making the default chat handler's BoundedChatHistory object globally accessible.

Fixes https://github.com/jupyterlab/jupyter-ai/issues/616 and https://github.com/jupyterlab/jupyter-ai/issues/817 (partially. fully with enhancement to /ask mentioned below)

Implementation summary

  • create a new global memory object 'llm_chat_history' which mirrors the 'chat_history' object
    • Move the BoundedChatHistory initiation to AiExtension.initialize_settings
    • Include it in 'self.settings' and as an attribute of BaseChatHandler
  • include traitlet 'default_max_chat_history' in AiExtension
  • add a line to ClearChatHandler to clear the chat memory.
  • changed BoundedChatHistory to not pop the internal memory object to fit size k but instead only return the last k*2 messages.
    • not directly related to features in this PR but will enable future features. Those involving dynamic configuration of history size and partial clearing (see Extensions)

Extensions

  • Reimplement AskChatHandler to use this llm_chat_history object.
    • Involves changing chain from ConversationBufferWindowMemory to RunnableWithMessageHistory.
    • ConversationBufferWindowMemory is deprecated and will need to be changed eventually anyway
    • I've left this for the team to implement as would involve some new prompt templates
  • FixChatHandler to add to llm_chat_history object.
    • May not be useful for /fix to use the chat history but may be useful to add to it.
    • This will allow user to ask follow up questions.
    • Some decisions need to be made regarding what to add as the human message.
    • Adding all the context may bloat the chat memory so a simple replacement to "Fix the errors in my notebook cell" + extra_instructions may suffice.
  • UI-based configuration of chat memory size by modifying llm_chat_history.k
  • Partial clearing on memory (reset to earlier chat state)
    • X button on user messages to remove all messages after
    • Allows users to reset to a previous chat history state
    • FYI, this is actually the feature that I want to implement and this PR is just a pre-requisite for that.

Possible amendments

Let me know if you would like me to rename:

  • memory object 'llm_chat_history'
  • traitlet 'default_max_chat_history' or the help description.
    • it is "default" to align with additional UI-based configuration in the future.

Although https://github.com/jupyterlab/jupyter-ai/issues/616 argues for making clearing of chat memory a separate command, https://github.com/jupyterlab/jupyter-ai/pull/842 seems to imply that has changed. Also, I found that most users assumed that /clear resets the memory perhaps because the /clear command in github copilot behaves that way. Let me know if you would like me to remove this from the PR.

I realise this is introducing multiple features and that this is all in a single PR. I mostly wanted to add the globally accessible memory but the rest were just natural and tiny additions. Let me know if you would like me to split it into 2 PRs instead (1. global memory + config, 2. /clear to reset memory)

Let me know what you guys think and if there is anything you would like me to add.

michaelchia avatar Aug 09 '24 09:08 michaelchia

@michaelchia Thanks for submitting this PR. Note that PR #842 was closed and not merged given that several modifications had been made to the base chat handler in the interim and this PR was going to be reworked from scratch as a new PR. You may assume PR 842 was not implemented, so this PR 943 supersedes 842. As you note, Issues #616 and #817 are related and it is good to handle them in one PR here.

Update the documentation in the section Additional chat commands in docs/source/users/index.md to explain how the change impacts both chat and memory histories, and how users should set history length k. image

srdas avatar Aug 12 '24 16:08 srdas

Note: Need to check if PR #938 which also moves the /clear handling to base.py has been taken into account in this PR.

srdas avatar Aug 15 '24 22:08 srdas

I don't see any conflict as the chat_memory clearing remained within ClearChatHandler in that PR. It's just the help message building and sending gets moved, which this PR doesn't touch. However, I wouldn't mind waiting for that to be approved first and merging its changes.

michaelchia avatar Aug 16 '24 01:08 michaelchia

@michaelchia - note that PR #938 is now merged. @dlqqq is away next week also, we may wait to make sure he does not have any further suggestions on this PR -- thanks for your detailed work and patience.

srdas avatar Aug 17 '24 07:08 srdas