Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix mypy
Browse files Browse the repository at this point in the history
dlqqq committed Dec 11, 2024
1 parent 9d15f22 commit cc467ec
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions packages/jupyter-ai/jupyter_ai/history.py
Original file line number Diff line number Diff line change
@@ -14,26 +14,32 @@

class YChatHistory(BaseChatMessageHistory):
"""
An implementation of `BaseChatMessageHistory` that yields the last `k`
An implementation of `BaseChatMessageHistory` that returns the preceding `k`
exchanges (`k * 2` messages) from the given YChat model.
If `k` is set to `None`, then this class returns all preceding messages.
"""

def __init__(self, ychat: YChat, k: Optional[int]):
self.ychat = ychat
self.k = k

@property
def messages(self) -> List[BaseMessage]:
"""Returns the last `k` messages."""
def messages(self) -> List[BaseMessage]: # type:ignore[override]
"""
Returns the last `2 * k` messages preceding the latest message. If
`k` is set to `None`, return all preceding messages.
"""
# TODO: consider bounding history based on message size (e.g. total
# char/token count) instead of message count.
all_messages = self.ychat.get_messages()

# gather last k * 2 messages and return
# we exclude the last message since that is the HumanMessage just
# we exclude the last message since that is the HumanChatMessage just
# submitted by a user.
messages = []
for message in all_messages[-self.k * 2 - 1 : -1]:
messages: List[BaseMessage] = []
start_idx = 0 if self.k is None else -2 * self.k - 1
for message in all_messages[start_idx:-1]:
if message["sender"] == BOT["username"]:
messages.append(AIMessage(content=message["body"]))
else:

0 comments on commit cc467ec

Please sign in to comment.