Skip to content

Commit

Permalink
Merge branch 'main' into jupyter-chat
Browse files Browse the repository at this point in the history
  • Loading branch information
brichet committed Oct 22, 2024
2 parents 2805efb + c2333d9 commit 7ea32a7
Show file tree
Hide file tree
Showing 13 changed files with 216 additions and 163 deletions.
39 changes: 37 additions & 2 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,43 @@

<!-- <START NEW CHANGELOG ENTRY> -->

## 2.26.0

This release notably includes the addition of a "Stop streaming" button, which takes over the "Send" button when a reply is streaming and the chat input is empty. While Jupyternaut is streaming a reply to a user, the user has the option to click the "Stop streaming" button to interrupt Jupyternaut and stop it from streaming further. Thank you @krassowski for contributing this feature! 🎉

([Full Changelog](https://github.com/jupyterlab/jupyter-ai/compare/@jupyter-ai/[email protected]))

### Enhancements made

- Support Quarto Markdown in `/learn` [#1047](https://github.com/jupyterlab/jupyter-ai/pull/1047) ([@dlqqq](https://github.com/dlqqq))
- Update requirements contributors doc [#1045](https://github.com/jupyterlab/jupyter-ai/pull/1045) ([@JasonWeill](https://github.com/JasonWeill))
- Remove clear_message_ids from RootChatHandler [#1042](https://github.com/jupyterlab/jupyter-ai/pull/1042) ([@michaelchia](https://github.com/michaelchia))
- Migrate streaming logic to `BaseChatHandler` [#1039](https://github.com/jupyterlab/jupyter-ai/pull/1039) ([@dlqqq](https://github.com/dlqqq))
- Unify message clearing & broadcast logic [#1038](https://github.com/jupyterlab/jupyter-ai/pull/1038) ([@dlqqq](https://github.com/dlqqq))
- Learn from JSON files [#1024](https://github.com/jupyterlab/jupyter-ai/pull/1024) ([@jlsajfj](https://github.com/jlsajfj))
- Allow users to stop message streaming [#1022](https://github.com/jupyterlab/jupyter-ai/pull/1022) ([@krassowski](https://github.com/krassowski))

### Bugs fixed

- Always use `username` from `IdentityProvider` [#1034](https://github.com/jupyterlab/jupyter-ai/pull/1034) ([@krassowski](https://github.com/krassowski))

### Maintenance and upkeep improvements

- Support `jupyter-collaboration` v3 [#1035](https://github.com/jupyterlab/jupyter-ai/pull/1035) ([@krassowski](https://github.com/krassowski))
- Test Python 3.9 and 3.12 on CI, test minimum dependencies [#1029](https://github.com/jupyterlab/jupyter-ai/pull/1029) ([@krassowski](https://github.com/krassowski))

### Documentation improvements

- Update requirements contributors doc [#1045](https://github.com/jupyterlab/jupyter-ai/pull/1045) ([@JasonWeill](https://github.com/JasonWeill))

### Contributors to this release

([GitHub contributors page for this release](https://github.com/jupyterlab/jupyter-ai/graphs/contributors?from=2024-10-07&to=2024-10-21&type=c))

[@dlqqq](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Adlqqq+updated%3A2024-10-07..2024-10-21&type=Issues) | [@JasonWeill](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3AJasonWeill+updated%3A2024-10-07..2024-10-21&type=Issues) | [@jlsajfj](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Ajlsajfj+updated%3A2024-10-07..2024-10-21&type=Issues) | [@krassowski](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Akrassowski+updated%3A2024-10-07..2024-10-21&type=Issues) | [@michaelchia](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Amichaelchia+updated%3A2024-10-07..2024-10-21&type=Issues) | [@pre-commit-ci](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Apre-commit-ci+updated%3A2024-10-07..2024-10-21&type=Issues)

<!-- <END NEW CHANGELOG ENTRY> -->

## 2.25.0

([Full Changelog](https://github.com/jupyterlab/jupyter-ai/compare/@jupyter-ai/[email protected]))
Expand All @@ -20,8 +57,6 @@

[@akaihola](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Aakaihola+updated%3A2024-10-04..2024-10-07&type=Issues) | [@dlqqq](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Adlqqq+updated%3A2024-10-04..2024-10-07&type=Issues) | [@jtpio](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Ajtpio+updated%3A2024-10-04..2024-10-07&type=Issues) | [@pre-commit-ci](https://github.com/search?q=repo%3Ajupyterlab%2Fjupyter-ai+involves%3Apre-commit-ci+updated%3A2024-10-04..2024-10-07&type=Issues)

<!-- <END NEW CHANGELOG ENTRY> -->

## 2.24.1

([Full Changelog](https://github.com/jupyterlab/jupyter-ai/compare/@jupyter-ai/[email protected]))
Expand Down
6 changes: 3 additions & 3 deletions docs/source/contributors/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,11 @@ Issues and pull requests that violate the above principles may be declined. If y

You can develop Jupyter AI on any system that can run a supported Python version up to and including 3.12, including recent Windows, macOS, and Linux versions.

Each Jupyter AI major version works with only one major version of JupyterLab. Jupyter AI 1.x supports JupyterLab 3.x, and Jupyter AI 2.x supports JupyterLab 4.x.
You should have the newest supported version of JupyterLab installed.

We highly recommend that you install [conda](https://conda.io/projects/conda/en/latest/user-guide/install/index.html) to start developing on Jupyter AI, especially if you are developing on macOS on an Apple Silicon-based Mac (M1, M1 Pro, M2, etc.).
We highly recommend that you install [conda](https://conda.io/projects/conda/en/latest/user-guide/install/index.html) to start contributing to Jupyter AI, especially if you are developing on macOS on an Apple Silicon-based Mac (M1, M1 Pro, M2, etc.).

You will need Node.js 18 to use Jupyter AI. Node.js 18.16.0 is known to work.
You will need [a supported version of node.js](https://github.com/nodejs/release#release-schedule) to use Jupyter AI.

:::{warning}
:name: node-18-15
Expand Down
2 changes: 1 addition & 1 deletion lerna.json
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
{
"$schema": "node_modules/lerna/schemas/lerna-schema.json",
"useWorkspaces": true,
"version": "2.25.0",
"version": "2.26.0",
"npmClient": "yarn",
"useNx": true
}
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "@jupyter-ai/monorepo",
"version": "2.25.0",
"version": "2.26.0",
"description": "A generative AI extension for JupyterLab",
"private": true,
"keywords": [
Expand Down
2 changes: 1 addition & 1 deletion packages/jupyter-ai-magics/package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "@jupyter-ai/magics",
"version": "2.25.0",
"version": "2.26.0",
"description": "Jupyter AI magics Python package. Not published on NPM.",
"private": true,
"homepage": "https://github.com/jupyterlab/jupyter-ai",
Expand Down
2 changes: 1 addition & 1 deletion packages/jupyter-ai-test/package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "@jupyter-ai/test",
"version": "2.25.0",
"version": "2.26.0",
"description": "Jupyter AI test package. Not published on NPM or PyPI.",
"private": true,
"homepage": "https://github.com/jupyterlab/jupyter-ai",
Expand Down
6 changes: 5 additions & 1 deletion packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,11 @@ async def process_message(self, message: HumanChatMessage, chat: Optional[YChat]
try:
with self.pending("Searching learned documents", message, chat=chat):
assert self.llm_chain
result = await self.llm_chain.acall({"question": query})
# TODO: migrate this class to use a LCEL `Runnable` instead of
# `Chain`, then remove the below ignore comment.
result = await self.llm_chain.acall( # type:ignore[attr-defined]
{"question": query}
)
response = result["answer"]
self.reply(response, chat, message)
except AssertionError as e:
Expand Down
163 changes: 151 additions & 12 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
import argparse
import asyncio
import contextlib
import os
import time
import traceback
from asyncio import Event
from typing import (
TYPE_CHECKING,
Any,
Awaitable,
Callable,
ClassVar,
Expand All @@ -21,10 +22,13 @@
from uuid import uuid4

from dask.distributed import Client as DaskClient
from jupyter_ai.callback_handlers import MetadataCallbackHandler
from jupyter_ai.config_manager import ConfigManager, Logger
from jupyter_ai.history import WrappedBoundedChatHistory
from jupyter_ai.models import (
AgentChatMessage,
AgentStreamChunkMessage,
AgentStreamMessage,
ChatMessage,
ClosePendingMessage,
HumanChatMessage,
Expand All @@ -33,8 +37,12 @@
)
from jupyter_ai_magics import Persona
from jupyter_ai_magics.providers import BaseProvider
from langchain.chains import LLMChain
from langchain.pydantic_v1 import BaseModel
from langchain_core.messages import AIMessageChunk
from langchain_core.runnables import Runnable
from langchain_core.runnables.config import RunnableConfig
from langchain_core.runnables.config import merge_configs as merge_runnable_configs
from langchain_core.runnables.utils import Input

try:
from jupyterlab_collaborative_chat.ychat import YChat
Expand Down Expand Up @@ -135,7 +143,7 @@ class BaseChatHandler:
"""Dictionary of context providers. Allows chat handlers to reference
context providers, which can be used to provide context to the LLM."""

message_interrupted: Dict[str, Event]
message_interrupted: Dict[str, asyncio.Event]
"""Dictionary mapping an agent message identifier to an asyncio Event
which indicates if the message generation/streaming was interrupted."""

Expand All @@ -153,7 +161,7 @@ def __init__(
help_message_template: str,
chat_handlers: Dict[str, "BaseChatHandler"],
context_providers: Dict[str, "BaseCommandContextProvider"],
message_interrupted: Dict[str, Event],
message_interrupted: Dict[str, asyncio.Event],
write_message: Callable[[YChat, str], None] | None = None,
):
self.log = log
Expand All @@ -180,7 +188,7 @@ def __init__(

self.llm: Optional[BaseProvider] = None
self.llm_params: Optional[dict] = None
self.llm_chain: Optional[LLMChain] = None
self.llm_chain: Optional[Runnable] = None

self.write_message = write_message

Expand Down Expand Up @@ -497,13 +505,6 @@ def send_help_message(
slash_commands_list=slash_commands_list,
context_commands_list=context_commands_list,
)
help_message = AgentChatMessage(
id=uuid4().hex,
time=time.time(),
body=help_message_body,
reply_to=human_msg.id if human_msg else "",
persona=self.persona,
)

if chat is not None:
self.write_message(chat, help_message_body)
Expand All @@ -516,3 +517,141 @@ def send_help_message(
persona=self.persona,
)
self.broadcast_message(help_message)

def _start_stream(self, human_msg: HumanChatMessage, chat: Optional[YChat]) -> str:
"""
Sends an `agent-stream` message to indicate the start of a response
stream. Returns the ID of the message, denoted as the `stream_id`.
"""
if chat is not None:
stream_id = self.write_message(chat, "")
else:
stream_id = uuid4().hex
stream_msg = AgentStreamMessage(
id=stream_id,
time=time.time(),
body="",
reply_to=human_msg.id,
persona=self.persona,
complete=False,
)

self.broadcast_message(stream_msg)

return stream_id

def _send_stream_chunk(
self,
stream_id: str,
content: str,
chat: Optional[YChat],
complete: bool = False,
metadata: Optional[Dict[str, Any]] = None,
) -> None:
"""
Sends an `agent-stream-chunk` message containing content that should be
appended to an existing `agent-stream` message with ID `stream_id`.
"""
if chat is not None:
self.write_message(chat, content, stream_id)
else:
if not metadata:
metadata = {}

stream_chunk_msg = AgentStreamChunkMessage(
id=stream_id, content=content, stream_complete=complete, metadata=metadata
)
self.broadcast_message(stream_chunk_msg)

async def stream_reply(
self,
input: Input,
human_msg: HumanChatMessage,
chat: Optional[YChat],
config: Optional[RunnableConfig] = None,
):
"""
Streams a reply to a human message by invoking
`self.llm_chain.astream()`. A LangChain `Runnable` instance must be
bound to `self.llm_chain` before invoking this method.
Arguments
---------
- `input`: The input to your runnable. The type of `input` depends on
the runnable in `self.llm_chain`, but is usually a dictionary whose keys
refer to input variables in your prompt template.
- `human_msg`: The `HumanChatMessage` being replied to.
- `config` (optional): A `RunnableConfig` object that specifies
additional configuration when streaming from the runnable.
"""
assert self.llm_chain
assert isinstance(self.llm_chain, Runnable)

received_first_chunk = False
metadata_handler = MetadataCallbackHandler()
base_config: RunnableConfig = {
"configurable": {"last_human_msg": human_msg},
"callbacks": [metadata_handler],
}
merged_config: RunnableConfig = merge_runnable_configs(base_config, config)

# start with a pending message
with self.pending("Generating response", human_msg, chat=chat) as pending_message:
# stream response in chunks. this works even if a provider does not
# implement streaming, as `astream()` defaults to yielding `_call()`
# when `_stream()` is not implemented on the LLM class.
chunk_generator = self.llm_chain.astream(input, config=merged_config)
stream_interrupted = False
async for chunk in chunk_generator:
if not received_first_chunk:
# when receiving the first chunk, close the pending message and
# start the stream.
self.close_pending(pending_message, chat=chat)
stream_id = self._start_stream(human_msg=human_msg, chat=chat)
received_first_chunk = True
self.message_interrupted[stream_id] = asyncio.Event()

if self.message_interrupted[stream_id].is_set():
try:
# notify the model provider that streaming was interrupted
# (this is essential to allow the model to stop generating)
#
# note: `mypy` flags this line, claiming that `athrow` is
# not defined on `AsyncIterator`. This is why an ignore
# comment is placed here.
await chunk_generator.athrow( # type:ignore[attr-defined]
GenerationInterrupted()
)
except GenerationInterrupted:
# do not let the exception bubble up in case if
# the provider did not handle it
pass
stream_interrupted = True
break

if isinstance(chunk, AIMessageChunk) and isinstance(chunk.content, str):
self._send_stream_chunk(stream_id, chunk.content, chat=chat)
elif isinstance(chunk, str):
self._send_stream_chunk(stream_id, chunk, chat=chat)
else:
self.log.error(f"Unrecognized type of chunk yielded: {type(chunk)}")
break

# complete stream after all chunks have been streamed
stream_tombstone = (
"\n\n(AI response stopped by user)" if stream_interrupted else ""
)
self._send_stream_chunk(
stream_id,
stream_tombstone,
chat=chat,
complete=True,
metadata=metadata_handler.jai_metadata,
)
del self.message_interrupted[stream_id]


class GenerationInterrupted(asyncio.CancelledError):
"""Exception raised when streaming is cancelled by the user"""
Loading

0 comments on commit 7ea32a7

Please sign in to comment.