|
1 | 1 | import argparse
|
2 |
| -from typing import Dict, Optional, Type |
| 2 | +from typing import Dict, Type |
3 | 3 |
|
4 | 4 | from jupyter_ai.models import HumanChatMessage
|
5 | 5 | from jupyter_ai_magics.providers import BaseProvider
|
6 |
| -from jupyterlab_chat.ychat import YChat |
7 | 6 | from langchain.chains import ConversationalRetrievalChain
|
8 | 7 | from langchain.memory import ConversationBufferWindowMemory
|
9 | 8 | from langchain_core.prompts import PromptTemplate
|
@@ -60,32 +59,32 @@ def create_llm_chain(
|
60 | 59 | verbose=False,
|
61 | 60 | )
|
62 | 61 |
|
63 |
| - async def process_message(self, message: HumanChatMessage, chat: Optional[YChat]): |
64 |
| - args = self.parse_args(message, chat) |
| 62 | + async def process_message(self, message: HumanChatMessage): |
| 63 | + args = self.parse_args(message) |
65 | 64 | if args is None:
|
66 | 65 | return
|
67 | 66 | query = " ".join(args.query)
|
68 | 67 | if not query:
|
69 |
| - self.reply(f"{self.parser.format_usage()}", chat, message) |
| 68 | + self.reply(f"{self.parser.format_usage()}", message) |
70 | 69 | return
|
71 | 70 |
|
72 | 71 | self.get_llm_chain()
|
73 | 72 |
|
74 | 73 | try:
|
75 |
| - with self.pending("Searching learned documents", message, chat=chat): |
| 74 | + with self.pending("Searching learned documents", message): |
76 | 75 | assert self.llm_chain
|
77 | 76 | # TODO: migrate this class to use a LCEL `Runnable` instead of
|
78 | 77 | # `Chain`, then remove the below ignore comment.
|
79 | 78 | result = await self.llm_chain.acall( # type:ignore[attr-defined]
|
80 | 79 | {"question": query}
|
81 | 80 | )
|
82 | 81 | response = result["answer"]
|
83 |
| - self.reply(response, chat, message) |
| 82 | + self.reply(response, message) |
84 | 83 | except AssertionError as e:
|
85 | 84 | self.log.error(e)
|
86 | 85 | response = """Sorry, an error occurred while reading the from the learned documents.
|
87 | 86 | If you have changed the embedding provider, try deleting the existing index by running
|
88 | 87 | `/learn -d` command and then re-submitting the `learn <directory>` to learn the documents,
|
89 | 88 | and then asking the question again.
|
90 | 89 | """
|
91 |
| - self.reply(response, chat, message) |
| 90 | + self.reply(response, message) |
0 commit comments