Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add metadata field to agent messages #1013

Merged
merged 5 commits into from
Sep 25, 2024
Merged

Conversation

dlqqq
Copy link
Member

@dlqqq dlqqq commented Sep 23, 2024

Description

Adds a metadata field to the agent message model, which should stores message metadata returned by the LLM provider. In LangChain, all providers have the option to return a generation_info dictionary when yielding one or more Generation objects while replying to a prompt. This PR makes that dictionary accessible via the metadata attribute in both the frontend & backend models for agent messages.

This dictionary may contain info about why an LLM finished generating a response (stop_reason) or append provider-specific data like a session_id. Some advanced custom use-cases like implementing thumbs up/down feedback in a custom message footer require this data to include in telemetry events.

In this implementation, this field is only set by DefaultChatHandler. For the time being, custom chat handlers have to mimic the implementation in this PR to include metadata in the returned message. In the future, the metadata field will be set automatically by BaseChatHandler so other chat handlers do not need to duplicate the implementation in their code. However, this requires refactoring DefaultChatHandler to migrate the streaming logic to BaseChatHandler.

Demo

Screen.Recording.2024-09-23.at.3.37.51.PM.mov

Testing

To replicate the above demo:

  1. Open packages/jupyter-ai/src/components/chat-messages.tsx.
  2. Add the below code to the top of ChatMessageHeader's definition to log the message metadata object:
export function ChatMessageHeader(props: ChatMessageHeaderProps): JSX.Element {
  const collaborators = useCollaboratorsContext();

  // new code here: vvv
  if (props.message.type === 'agent-stream' && props.message.complete) {
    console.log(props.message.metadata);
  }
  
  ...
}
  1. Run jlpm build and restart JupyterLab.
  2. Verify that message metadata is logged in the console after receiving a message. You can use test-provider-with-streaming:test model provided by jupyter_ai_test (installable via jlpm dev-uninstall && jlpm dev-install if not already installed). It should always set the message metadata to {"test_metadata_field":"foobar"}.

@dlqqq dlqqq added the enhancement New feature or request label Sep 23, 2024
@srdas
Copy link
Collaborator

srdas commented Sep 24, 2024

I am seeing the {"test_metadata_field":"foobar"} as expected:
image

Replicated the same with another LLM:
image

However, with Bedrock Claude3-haiku there is no metadata returned (same for Anthropic Claude3-haiku):
image

@srdas srdas self-requested a review September 24, 2024 22:38
Copy link
Collaborator

@srdas srdas left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See testing comments in the conversation.

@dlqqq dlqqq merged commit 0884211 into jupyterlab:main Sep 25, 2024
9 checks passed
@dlqqq dlqqq deleted the add-message-metadata branch September 25, 2024 17:33
dlqqq added a commit to dlqqq/jupyter-ai that referenced this pull request Oct 2, 2024
dlqqq added a commit to dlqqq/jupyter-ai that referenced this pull request Oct 7, 2024
dlqqq added a commit to dlqqq/jupyter-ai that referenced this pull request Oct 7, 2024
dlqqq added a commit to dlqqq/jupyter-ai that referenced this pull request Oct 7, 2024
dlqqq added a commit that referenced this pull request Oct 7, 2024
* remove console log accidentally merged with #1013

* set metadata on stream messages in the chat_history array

* implement support for optional telemetry plugin

* anonymize message & code details

* export telemetry hook from NPM package entry point
Marchlak pushed a commit to Marchlak/jupyter-ai that referenced this pull request Oct 28, 2024
* add `metadata` field to agent messages

Currently only set when the default chat handler is used.

* include message metadata in TestLLMWithStreaming

* tweak TestLLMWithStreaming parameters

* pre-commit

* default to empty dict if no generation_info
Marchlak pushed a commit to Marchlak/jupyter-ai that referenced this pull request Oct 28, 2024
* remove console log accidentally merged with jupyterlab#1013

* set metadata on stream messages in the chat_history array

* implement support for optional telemetry plugin

* anonymize message & code details

* export telemetry hook from NPM package entry point
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants