From cc467ec68a95660ab27edb6d39e0e5babd22fc4a Mon Sep 17 00:00:00 2001 From: "David L. Qiu" Date: Wed, 11 Dec 2024 12:31:39 -0800 Subject: [PATCH] fix mypy --- packages/jupyter-ai/jupyter_ai/history.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/history.py b/packages/jupyter-ai/jupyter_ai/history.py index fbc7d5934..b3a7b8165 100644 --- a/packages/jupyter-ai/jupyter_ai/history.py +++ b/packages/jupyter-ai/jupyter_ai/history.py @@ -14,8 +14,10 @@ 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]): @@ -23,17 +25,21 @@ def __init__(self, ychat: YChat, k: Optional[int]): 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: