diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 0702d0e6f..f7bb8f085 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "ipython", "pydantic", "importlib_metadata~=5.2.0", - "langchain~=0.0.115" + "langchain~=0.0.144" ] [project.optional-dependencies] diff --git a/packages/jupyter-ai/jupyter_ai/actors/filesystem.py b/packages/jupyter-ai/jupyter_ai/actors/ask.py similarity index 57% rename from packages/jupyter-ai/jupyter_ai/actors/filesystem.py rename to packages/jupyter-ai/jupyter_ai/actors/ask.py index 7d48609b2..5676154f0 100644 --- a/packages/jupyter-ai/jupyter_ai/actors/filesystem.py +++ b/packages/jupyter-ai/jupyter_ai/actors/ask.py @@ -1,16 +1,18 @@ -from langchain import OpenAI +import argparse + import ray -import time -from uuid import uuid4 from ray.util.queue import Queue + +from langchain import OpenAI from langchain.chains import ConversationalRetrievalChain -from jupyter_ai.models import AgentChatMessage, HumanChatMessage + +from jupyter_ai.models import HumanChatMessage from jupyter_ai.actors.base import ACTOR_TYPE, BaseActor, Logger @ray.remote -class FileSystemActor(BaseActor): - """Processes messages prefixed with /fs. This actor will +class AskActor(BaseActor): + """Processes messages prefixed with /ask. This actor will send the message as input to a RetrieverQA chain, that follows the Retrieval and Generation (RAG) tehnique to query the documents from the index, and sends this context @@ -18,8 +20,8 @@ class FileSystemActor(BaseActor): """ def __init__(self, reply_queue: Queue, log: Logger): - super().__init__(log=log, reply_queue=reply_queue) - index_actor = ray.get_actor(ACTOR_TYPE.INDEX.value) + super().__init__(reply_queue=reply_queue, log=log) + index_actor = ray.get_actor(ACTOR_TYPE.LEARN.value) handle = index_actor.get_index.remote() vectorstore = ray.get(handle) if not vectorstore: @@ -31,22 +33,26 @@ def __init__(self, reply_queue: Queue, log: Logger): vectorstore.as_retriever() ) - def process_message(self, message: HumanChatMessage): - query = message.body.split(' ', 1)[-1] + self.parser.prog = '/ask' + self.parser.add_argument('query', nargs=argparse.REMAINDER) + + + def _process_message(self, message: HumanChatMessage): + args = self.parse_args(message) + if args is None: + return + query = ' '.join(args.query) + if not query: + self.reply(f"{self.parser.format_usage()}", message) + return - index_actor = ray.get_actor(ACTOR_TYPE.INDEX.value) + index_actor = ray.get_actor(ACTOR_TYPE.LEARN.value) handle = index_actor.get_index.remote() vectorstore = ray.get(handle) # Have to reference the latest index self.chat_provider.retriever = vectorstore.as_retriever() result = self.chat_provider({"question": query, "chat_history": self.chat_history}) - reply = result['answer'] - self.chat_history.append((query, reply)) - agent_message = AgentChatMessage( - id=uuid4().hex, - time=time.time(), - body=reply, - reply_to=message.id - ) - self.reply_queue.put(agent_message) \ No newline at end of file + response = result['answer'] + self.chat_history.append((query, response)) + self.reply(response, message) diff --git a/packages/jupyter-ai/jupyter_ai/actors/base.py b/packages/jupyter-ai/jupyter_ai/actors/base.py index 5d79d508a..3cd9318a4 100644 --- a/packages/jupyter-ai/jupyter_ai/actors/base.py +++ b/packages/jupyter-ai/jupyter_ai/actors/base.py @@ -1,22 +1,27 @@ +import argparse from enum import Enum +from uuid import uuid4 +import time import logging from typing import Union -from jupyter_ai.models import HumanChatMessage +import traceback + from ray.util.queue import Queue +from jupyter_ai.models import HumanChatMessage, AgentChatMessage + Logger = Union[logging.Logger, logging.LoggerAdapter] class ACTOR_TYPE(str, Enum): DEFAULT = "default" - FILESYSTEM = "filesystem" - INDEX = 'index' + ASK = "ask" + LEARN = 'learn' MEMORY = 'memory' COMMANDS = { - '/fs': ACTOR_TYPE.FILESYSTEM, - '/filesystem': ACTOR_TYPE.FILESYSTEM, - '/index': ACTOR_TYPE.INDEX + '/ask': ACTOR_TYPE.ASK, + '/learn': ACTOR_TYPE.LEARN } class BaseActor(): @@ -29,7 +34,36 @@ def __init__( ): self.log = log self.reply_queue = reply_queue + self.parser = argparse.ArgumentParser() def process_message(self, message: HumanChatMessage): """Processes the message passed by the `Router`""" - raise NotImplementedError("Should be implemented by subclasses.") \ No newline at end of file + try: + self._process_message(message) + except Exception as e: + formatted_e = traceback.format_exc() + response = f"Sorry, something went wrong and I wasn't able to index that path.\n\n```\n{formatted_e}\n```" + self.reply(response, message) + + def _process_message(self, message: HumanChatMessage): + """Processes the message passed by the `Router`""" + raise NotImplementedError("Should be implemented by subclasses.") + + def reply(self, response, message: HumanChatMessage): + m = AgentChatMessage( + id=uuid4().hex, + time=time.time(), + body=response, + reply_to=message.id + ) + self.reply_queue.put(m) + + def parse_args(self, message): + args = message.body.split(' ') + try: + args = self.parser.parse_args(args[1:]) + except (argparse.ArgumentError, SystemExit) as e: + response = f"{self.parser.format_usage()}" + self.reply(response, message) + return None + return args \ No newline at end of file diff --git a/packages/jupyter-ai/jupyter_ai/actors/default.py b/packages/jupyter-ai/jupyter_ai/actors/default.py index 7bfb76873..6ab33ae08 100644 --- a/packages/jupyter-ai/jupyter_ai/actors/default.py +++ b/packages/jupyter-ai/jupyter_ai/actors/default.py @@ -1,13 +1,7 @@ -import time -from uuid import uuid4 -from jupyter_ai.actors.base import BaseActor, Logger, ACTOR_TYPE -from jupyter_ai.actors.memory import RemoteMemory -from jupyter_ai.models import AgentChatMessage, HumanChatMessage -from jupyter_ai_magics.providers import ChatOpenAINewProvider -from langchain import ConversationChain import ray from ray.util.queue import Queue -from langchain.memory import ConversationBufferMemory + +from langchain import ConversationChain from langchain.prompts import ( ChatPromptTemplate, MessagesPlaceholder, @@ -15,14 +9,17 @@ HumanMessagePromptTemplate ) -SYSTEM_PROMPT = "The following is a friendly conversation between a human and an AI, whose name is Jupyter AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know." +from jupyter_ai.actors.base import BaseActor, Logger, ACTOR_TYPE +from jupyter_ai.actors.memory import RemoteMemory +from jupyter_ai.models import HumanChatMessage +from jupyter_ai_magics.providers import ChatOpenAINewProvider +SYSTEM_PROMPT = "The following is a friendly conversation between a human and an AI, whose name is Jupyter AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know." @ray.remote class DefaultActor(BaseActor): def __init__(self, reply_queue: Queue, log: Logger): - super().__init__(log=log, reply_queue=reply_queue) - # TODO: Should take the provider/model id as strings + super().__init__(reply_queue=reply_queue, log=log) provider = ChatOpenAINewProvider(model_id="gpt-3.5-turbo") # Create a conversation memory @@ -40,12 +37,6 @@ def __init__(self, reply_queue: Queue, log: Logger): ) self.chat_provider = chain - def process_message(self, message: HumanChatMessage): + def _process_message(self, message: HumanChatMessage): response = self.chat_provider.predict(input=message.body) - agent_message = AgentChatMessage( - id=uuid4().hex, - time=time.time(), - body=response, - reply_to=message.id - ) - self.reply_queue.put(agent_message) \ No newline at end of file + self.reply(response, message) diff --git a/packages/jupyter-ai/jupyter_ai/actors/index.py b/packages/jupyter-ai/jupyter_ai/actors/index.py deleted file mode 100644 index 3349d0f59..000000000 --- a/packages/jupyter-ai/jupyter_ai/actors/index.py +++ /dev/null @@ -1,63 +0,0 @@ -import time -from uuid import uuid4 -from jupyter_ai.models import AgentChatMessage, HumanChatMessage -from langchain import FAISS -import ray -import os - -from jupyter_ai.actors.base import BaseActor, Logger -from jupyter_ai_magics.providers import ChatOpenAINewProvider - -from ray.util.queue import Queue -from jupyter_core.paths import jupyter_data_dir -from langchain.embeddings.openai import OpenAIEmbeddings -from langchain.document_loaders import TextLoader -from langchain.text_splitter import RecursiveCharacterTextSplitter -from jupyter_ai.document_loaders.directory import DirectoryLoader - - -@ray.remote -class DocumentIndexActor(BaseActor): - def __init__(self, reply_queue: Queue, root_dir: str, log: Logger): - super().__init__(log=log, reply_queue=reply_queue) - self.root_dir = root_dir - self.index_save_dir = os.path.join(jupyter_data_dir(), '.jupyter_ai_indexes') - - if ChatOpenAINewProvider.auth_strategy.name not in os.environ: - return - - if not os.path.exists(self.index_save_dir): - os.makedirs(self.index_save_dir) - - embeddings = OpenAIEmbeddings() - try: - self.index = FAISS.load_local(self.index_save_dir, embeddings) - except Exception as e: - self.index = FAISS.from_texts(["This is just starter text for the index"], embeddings) - - def get_index(self): - return self.index - - def process_message(self, message: HumanChatMessage): - dir_path = message.body.split(' ', 1)[-1] - load_path = os.path.join(self.root_dir, dir_path) - loader = DirectoryLoader( - load_path, - glob=['*.txt', '*.text', '*.md', '*.py', '*.ipynb', '*.R'], - loader_cls=TextLoader - ) - documents = loader.load_and_split( - text_splitter=RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20) - ) - self.index.add_documents(documents) - self.index.save_local(self.index_save_dir) - - response = f"""🎉 I have indexed documents at **{dir_path}** and ready to answer questions. - You can ask questions from these docs by prefixing your message with **/fs**.""" - agent_message = AgentChatMessage( - id=uuid4().hex, - time=time.time(), - body=response, - reply_to=message.id - ) - self.reply_queue.put(agent_message) \ No newline at end of file diff --git a/packages/jupyter-ai/jupyter_ai/actors/learn.py b/packages/jupyter-ai/jupyter_ai/actors/learn.py new file mode 100644 index 000000000..c32c1b054 --- /dev/null +++ b/packages/jupyter-ai/jupyter_ai/actors/learn.py @@ -0,0 +1,119 @@ +import os +import traceback +from collections import Counter +import argparse + +import ray +from ray.util.queue import Queue + +from jupyter_core.paths import jupyter_data_dir + +from langchain import FAISS +from langchain.embeddings.openai import OpenAIEmbeddings +from langchain.text_splitter import ( + RecursiveCharacterTextSplitter, PythonCodeTextSplitter, + MarkdownTextSplitter, LatexTextSplitter +) + +from jupyter_ai.models import HumanChatMessage +from jupyter_ai.actors.base import BaseActor, Logger +from jupyter_ai_magics.providers import ChatOpenAINewProvider +from jupyter_ai.document_loaders.directory import RayRecursiveDirectoryLoader +from jupyter_ai.document_loaders.splitter import ExtensionSplitter, NotebookSplitter + + +@ray.remote +class LearnActor(BaseActor): + + def __init__(self, reply_queue: Queue, log: Logger, root_dir: str): + super().__init__(reply_queue=reply_queue, log=log) + self.root_dir = root_dir + self.index_save_dir = os.path.join(jupyter_data_dir(), 'jupyter_ai', 'indices') + self.chunk_size = 2000 + self.chunk_overlap = 100 + self.parser.prog = '/learn' + self.parser.add_argument('-v', '--verbose', action='store_true') + self.parser.add_argument('-d', '--delete', action='store_true') + self.parser.add_argument('path', nargs=argparse.REMAINDER) + self.index_name = 'default' + self.index = None + + if ChatOpenAINewProvider.auth_strategy.name not in os.environ: + return + + if not os.path.exists(self.index_save_dir): + os.makedirs(self.index_save_dir) + + self.load_or_create() + + def _process_message(self, message: HumanChatMessage): + args = self.parse_args(message) + if args is None: + return + + if args.delete: + self.delete() + self.reply(f"👍 I have deleted everything I previously learned.", message) + return + + # Make sure the path exists. + if not len(args.path) == 1: + self.reply(f"{self.parser.format_usage()}", message) + return + short_path = args.path[0] + load_path = os.path.join(self.root_dir, short_path) + if not os.path.exists(load_path): + response = f"Sorry, that path doesn't exist: {load_path}" + self.reply(response, message) + return + + if args.verbose: + self.reply(f"Loading and splitting files for {load_path}", message) + + splitters={ + '.py': PythonCodeTextSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap), + '.md': MarkdownTextSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap), + '.tex': LatexTextSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap), + '.ipynb': NotebookSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap) + } + splitter = ExtensionSplitter( + splitters=splitters, + default_splitter=RecursiveCharacterTextSplitter(chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap) + ) + + loader = RayRecursiveDirectoryLoader(load_path) + texts = loader.load_and_split(text_splitter=splitter) + self.index.add_documents(texts) + self.save() + + response = f"""🎉 I have indexed documents at **{load_path}** and I am ready to answer questions about them. + You can ask questions about these docs by prefixing your message with **/ask**.""" + self.reply(response, message) + + def get_index(self): + return self.index + + def delete(self): + self.index = None + paths = [os.path.join(self.index_save_dir, self.index_name+ext) for ext in ['.pkl', '.faiss']] + for path in paths: + if os.path.isfile(path): + os.remove(path) + self.create() + + def create(self): + embeddings = OpenAIEmbeddings() + self.index = FAISS.from_texts(["Jupyter AI knows about your filesystem, to ask questions first use the /learn command."], embeddings) + self.save() + + def save(self): + if self.index is not None: + self.index.save_local(self.index_save_dir, index_name=self.index_name) + + def load_or_create(self): + embeddings = OpenAIEmbeddings() + if self.index is None: + try: + self.index = FAISS.load_local(self.index_save_dir, embeddings, index_name=self.index_name) + except Exception as e: + self.create() diff --git a/packages/jupyter-ai/jupyter_ai/actors/memory.py b/packages/jupyter-ai/jupyter_ai/actors/memory.py index 6271aa943..808537599 100644 --- a/packages/jupyter-ai/jupyter_ai/actors/memory.py +++ b/packages/jupyter-ai/jupyter_ai/actors/memory.py @@ -1,9 +1,10 @@ -from jupyter_ai.actors.base import Logger from typing import Dict, Any, List -from langchain.schema import BaseMemory + import ray +from langchain.schema import BaseMemory from pydantic import PrivateAttr +from jupyter_ai.actors.base import Logger @ray.remote class MemoryActor(object): @@ -13,7 +14,7 @@ class MemoryActor(object): running in different actors by using RemoteMemory (below). """ - def __init__(self, memory: BaseMemory, log: Logger): + def __init__(self, log: Logger, memory: BaseMemory): self.memory = memory self.log = log diff --git a/packages/jupyter-ai/jupyter_ai/actors/router.py b/packages/jupyter-ai/jupyter_ai/actors/router.py index 72bbd1268..7b417a0cd 100644 --- a/packages/jupyter-ai/jupyter_ai/actors/router.py +++ b/packages/jupyter-ai/jupyter_ai/actors/router.py @@ -1,9 +1,8 @@ -from jupyter_ai.actors.base import ACTOR_TYPE, COMMANDS, Logger, BaseActor -from jupyter_ai.models import ClearMessage - import ray from ray.util.queue import Queue +from jupyter_ai.actors.base import ACTOR_TYPE, COMMANDS, Logger, BaseActor +from jupyter_ai.models import ClearMessage @ray.remote class Router(BaseActor): @@ -13,9 +12,9 @@ def __init__(self, reply_queue: Queue, log: Logger): To register new actors, add the actor type in the `ACTOR_TYPE` enum and add a corresponding command in the `COMMANDS` dictionary. """ - super().__init__(log=log, reply_queue=reply_queue) + super().__init__(reply_queue=reply_queue, log=log) - def route_message(self, message): + def _process_message(self, message): # assign default actor default = ray.get_actor(ACTOR_TYPE.DEFAULT) @@ -30,4 +29,3 @@ def route_message(self, message): self.reply_queue.put(reply_message) else: default.process_message.remote(message) - diff --git a/packages/jupyter-ai/jupyter_ai/document_loaders/directory.py b/packages/jupyter-ai/jupyter_ai/document_loaders/directory.py index ff380dfdb..f87a74e30 100644 --- a/packages/jupyter-ai/jupyter_ai/document_loaders/directory.py +++ b/packages/jupyter-ai/jupyter_ai/document_loaders/directory.py @@ -1,54 +1,73 @@ -import logging -from typing import Dict, List, Optional, Union +from typing import List, Optional +from pathlib import Path +import hashlib +import itertools + +import ray from langchain.document_loaders.base import BaseLoader -from langchain.document_loaders.directory import FILE_LOADER_TYPE, _is_visible -from langchain.document_loaders.unstructured import UnstructuredFileLoader from langchain.schema import Document -from wcmatch.pathlib import Path - - -logger = logging.getLogger(__name__) +from langchain.document_loaders.directory import _is_visible +from langchain.text_splitter import ( + RecursiveCharacterTextSplitter, TextSplitter, +) +@ray.remote +def path_to_doc(path): + with open(str(path), 'r') as f: + text = f.read() + m = hashlib.sha256() + m.update(text.encode('utf-8')) + metadata = {'path': str(path), 'sha256': m.digest(), 'extension': path.suffix} + return Document(page_content=text, metadata=metadata) -class DirectoryLoader(BaseLoader): - """Loading logic for loading documents from a directory.""" +class ExcludePattern(Exception): + pass + +def iter_paths(path, extensions, exclude): + for p in Path(path).rglob('*'): + if p.is_dir(): + continue + if not _is_visible(p.relative_to(path)): + continue + try: + for pattern in exclude: + if pattern in str(p): + raise ExcludePattern() + except ExcludePattern: + continue + if p.suffix in extensions: + yield p +class RayRecursiveDirectoryLoader(BaseLoader): + def __init__( self, - path: str, - glob: Union[str, List[str]] = "**/[!.]*", - silent_errors: bool = False, - load_hidden: bool = False, - loader_cls: FILE_LOADER_TYPE = UnstructuredFileLoader, - loader_kwargs: Optional[Dict] = None, - recursive: bool = False, + path, + extensions={'.py', '.md', '.R', '.Rmd', '.jl', '.sh', '.ipynb', '.js', '.ts', '.jsx', '.tsx', '.txt'}, + exclude={'.ipynb_checkpoints', 'node_modules', 'lib', 'build', '.git', '.DS_Store'} ): - """Initialize with path to directory and how to glob over it.""" - if loader_kwargs is None: - loader_kwargs = {} self.path = path - self.glob = glob - self.load_hidden = load_hidden - self.loader_cls = loader_cls - self.loader_kwargs = loader_kwargs - self.silent_errors = silent_errors - self.recursive = recursive - + self.extensions = extensions + self.exclude=exclude + def load(self) -> List[Document]: - """Load documents.""" - p = Path(self.path) - docs = [] - items = p.rglob(self.glob) if self.recursive else p.glob(self.glob) - for i in items: - if i.is_file(): - if _is_visible(i.relative_to(p)) or self.load_hidden: - try: - sub_docs = self.loader_cls(str(i), **self.loader_kwargs).load() - docs.extend(sub_docs) - except Exception as e: - if self.silent_errors: - logger.warning(e) - else: - raise e - return docs \ No newline at end of file + paths = iter_paths(self.path, self.extensions, self.exclude) + doc_refs = list(map(path_to_doc.remote, paths)) + return ray.get(doc_refs) + + def load_and_split( + self, text_splitter: Optional[TextSplitter] = None + ) -> List[Document]: + if text_splitter is None: + _text_splitter = RecursiveCharacterTextSplitter() + else: + _text_splitter = text_splitter + + @ray.remote + def split(doc): + return _text_splitter.split_documents([doc]) + + paths = iter_paths(self.path, self.extensions, self.exclude) + doc_refs = map(split.remote, map(path_to_doc.remote, paths)) + return list(itertools.chain(*ray.get(list(doc_refs)))) \ No newline at end of file diff --git a/packages/jupyter-ai/jupyter_ai/document_loaders/splitter.py b/packages/jupyter-ai/jupyter_ai/document_loaders/splitter.py new file mode 100644 index 000000000..d80289453 --- /dev/null +++ b/packages/jupyter-ai/jupyter_ai/document_loaders/splitter.py @@ -0,0 +1,44 @@ +from typing import List, Optional + +from langchain.schema import Document +from langchain.text_splitter import TextSplitter, RecursiveCharacterTextSplitter, MarkdownTextSplitter +import copy + +class ExtensionSplitter(TextSplitter): + + def __init__(self, splitters, default_splitter=None): + self.splitters = splitters + if default_splitter is None: + self.default_splitter = RecursiveCharacterTextSplitter() + else: + self.default_splitter = default_splitter + + def split_text(self, text: str, metadata=None): + splitter = self.splitters.get(metadata['extension'], self.default_splitter) + return splitter.split_text(text) + + def create_documents(self, texts: List[str], metadatas: Optional[List[dict]] = None) -> List[Document]: + _metadatas = metadatas or [{}] * len(texts) + documents = [] + for i, text in enumerate(texts): + metadata = copy.deepcopy(_metadatas[i]) + for chunk in self.split_text(text, metadata): + new_doc = Document( + page_content=chunk, metadata=metadata + ) + documents.append(new_doc) + return documents + +import nbformat + +class NotebookSplitter(TextSplitter): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.markdown_splitter = MarkdownTextSplitter(chunk_size=self._chunk_size, chunk_overlap=self._chunk_overlap) + + def split_text(self, text: str): + nb = nbformat.reads(text, as_version=4) + md = '\n\n'.join([cell.source for cell in nb.cells]) + return self.markdown_splitter.split_text(md) + diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index a4c1ff0b0..a72988636 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -3,8 +3,8 @@ import queue from langchain.memory import ConversationBufferWindowMemory from jupyter_ai.actors.default import DefaultActor -from jupyter_ai.actors.filesystem import FileSystemActor -from jupyter_ai.actors.index import DocumentIndexActor +from jupyter_ai.actors.ask import AskActor +from jupyter_ai.actors.learn import LearnActor from jupyter_ai.actors.router import Router from jupyter_ai.actors.memory import MemoryActor from jupyter_ai.actors.base import ACTOR_TYPE @@ -120,23 +120,23 @@ def initialize_settings(self): reply_queue=reply_queue, log=self.log ) - index_actor = DocumentIndexActor.options(name=ACTOR_TYPE.INDEX.value).remote( + learn_actor = LearnActor.options(name=ACTOR_TYPE.LEARN.value).remote( reply_queue=reply_queue, + log=self.log, root_dir=self.serverapp.root_dir, - log=self.log ) - fs_actor = FileSystemActor.options(name=ACTOR_TYPE.FILESYSTEM.value).remote( + ask_actor = AskActor.options(name=ACTOR_TYPE.ASK.value).remote( reply_queue=reply_queue, log=self.log ) memory_actor = MemoryActor.options(name=ACTOR_TYPE.MEMORY.value).remote( - memory=ConversationBufferWindowMemory(return_messages=True, k=2), - log=self.log + log=self.log, + memory=ConversationBufferWindowMemory(return_messages=True, k=2) ) self.settings['router'] = router self.settings["default_actor"] = default_actor - self.settings["index_actor"] = index_actor - self.settings["fs_actor"] = fs_actor + self.settings["learn_actor"] = learn_actor + self.settings["ask_actor"] = ask_actor self.settings["memory_actor"] = memory_actor reply_processor = ReplyProcessor(self.settings['chat_handlers'], reply_queue, log=self.log) diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index 30e1e78d0..fa2ad1dc6 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -228,7 +228,7 @@ async def on_message(self, message): # process through the router router = ray.get_actor("router") - router.route_message.remote(chat_message) + router.process_message.remote(chat_message) def on_close(self): self.log.debug("Disconnecting client with user %s", self.client_id) diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml index 57bdc4b42..6e4ada2a6 100644 --- a/packages/jupyter-ai/pyproject.toml +++ b/packages/jupyter-ai/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "openai~=0.26", "aiosqlite~=0.18", "importlib_metadata~=5.2.0", - "langchain~=0.0.115", + "langchain~=0.0.144", "tiktoken", # required for OpenAIEmbeddings "jupyter_ai_magics", "ray==2.2.0", # latest ray version 2.3.0 requires grpcio installation from conda