Moved actors to separate modules.

packages/jupyter-ai/jupyter_ai/actors/__init__.py

@@ -0,0 +1 @@
+"""Actor classes that process incoming chat messages"""

packages/jupyter-ai/jupyter_ai/actors/base.py

@@ -0,0 +1,34 @@
+from enum import Enum
+import logging
+from typing import Union
+from jupyter_ai.models import HumanChatMessage
+from ray.util.queue import Queue
+
+
+Logger = Union[logging.Logger, logging.LoggerAdapter]
+
+class ACTOR_TYPE(str, Enum):
+    DEFAULT = "default"
+    FILESYSTEM = "filesystem"
+    INDEX = 'index'
+
+COMMANDS = {
+    '/fs': ACTOR_TYPE.FILESYSTEM,
+    '/filesystem': ACTOR_TYPE.FILESYSTEM,
+    '/index': ACTOR_TYPE.INDEX
+}
+
+class BaseActor():
+    """Base actor implemented by actors that are called by the `Router`"""
+
+    def __init__(
+            self, 
+            log: Logger,
+            reply_queue: Queue
+        ):
+        self.log = log
+        self.reply_queue = reply_queue
+
+    def process_message(self, message: HumanChatMessage):
+        """Processes the message passed by the `Router`"""
+        raise NotImplementedError("Should be implemented by subclasses.")

packages/jupyter-ai/jupyter_ai/actors/default.py

@@ -0,0 +1,49 @@
+import time
+from uuid import uuid4
+from jupyter_ai.actors.base import BaseActor, Logger
+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.prompts import (
+    ChatPromptTemplate, 
+    MessagesPlaceholder, 
+    SystemMessagePromptTemplate, 
+    HumanMessagePromptTemplate
+)
+
+
+
+@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
+        provider = ChatOpenAINewProvider(model_id="gpt-3.5-turbo")
+
+        # Create a conversation memory
+        memory = ConversationBufferMemory(return_messages=True)
+        prompt_template = ChatPromptTemplate.from_messages([
+            SystemMessagePromptTemplate.from_template("The following is a friendly conversation between a human and an 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."),
+            MessagesPlaceholder(variable_name="history"),
+            HumanMessagePromptTemplate.from_template("{input}")
+        ])
+        chain = ConversationChain(
+            llm=provider, 
+            prompt=prompt_template,
+            verbose=True, 
+            memory=memory
+        )
+        self.chat_provider = chain
+
+    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)

packages/jupyter-ai/jupyter_ai/actors/filesystem.py

@@ -0,0 +1,52 @@
+from langchain import OpenAI
+import ray
+import time
+from uuid import uuid4
+from ray.util.queue import Queue
+from langchain.chains import ConversationalRetrievalChain
+from jupyter_ai.models import AgentChatMessage, HumanChatMessage
+from jupyter_ai.actors.base import ACTOR_TYPE, BaseActor, Logger
+
+
+@ray.remote
+class FileSystemActor(BaseActor):
+    """Processes messages prefixed with /fs. 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
+    to the LLM to generate the final reply.
+    """
+
+    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)
+        handle = index_actor.get_index.remote()
+        vectorstore = ray.get(handle)
+        if not vectorstore:
+            return
+
+        self.chat_history = []
+        self.chat_provider = ConversationalRetrievalChain.from_llm(
+            OpenAI(temperature=0, verbose=True),
+            vectorstore.as_retriever()
+        )
+
+    def process_message(self, message: HumanChatMessage):
+        query = message.body.split(' ', 1)[-1]
+
+        index_actor = ray.get_actor(ACTOR_TYPE.INDEX.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)