Some refactoring the memory inerface

megabots/__init__.py

@@ -1,256 +1,9 @@
-from typing import Any
-from langchain.llms import OpenAI
-from langchain.chat_models import ChatOpenAI
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.chains.qa_with_sources import load_qa_with_sources_chain
-from langchain.vectorstores.faiss import FAISS
-import gradio as gr
-from fastapi import FastAPI
-import pickle
-import os
-from dotenv import load_dotenv
-from langchain.prompts import PromptTemplate
-from langchain.chains.question_answering import load_qa_chain
-from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
-from langchain.document_loaders import DirectoryLoader
-from megabots.vectorstores import VectorStore, vectorstore
-from langchain.memory import (
-    ConversationBufferMemory,
-    ConversationBufferWindowMemory,
-    ConversationSummaryMemory,
-    ConversationSummaryBufferMemory,
-)
-
-load_dotenv()
-
-
-class Bot:
-    def __init__(
-        self,
-        model: str | None = None,
-        prompt_template: str | None = None,
-        prompt_variables: list[str] | None = None,
-        index: str | None = None,
-        sources: bool | None = False,
-        vectorstore: VectorStore | None = None,
-        memory: str | None = None,
-        memory_window: int = 3,
-        verbose: bool = False,
-        temperature: int = 0,
-    ):
-        self.select_model(model, temperature)
-        self.create_loader(index)
-        self.load_or_create_index(index, vectorstore)
-
-        # Load the question-answering chain for the selected model
-        self.chain = self.create_chain(
-            prompt_template, prompt_variables, sources=sources, verbose=verbose
-        )
-
- 
-
-    def create_chain(
-        self,
-        prompt_template: str | None = None,
-        prompt_variables: list[str] | None = None,
-        sources: bool | None = False,
-        verbose: bool = False,
-    ):
-        prompt = (
-            PromptTemplate(template=prompt_template, input_variables=prompt_variables)
-            if prompt_template is not None and prompt_variables is not None
-            else QA_PROMPT
-        )
-        # TODO: Changing the prompt here is not working. Leave it as is for now.
-        # Reference: https://github.com/hwchase17/langchain/issues/2858
-        if sources:
-            return load_qa_with_sources_chain(
-                self.llm, chain_type="stuff", verbose=verbose
-            )
-        return load_qa_chain(
-            self.llm, chain_type="stuff", verbose=verbose, prompt=prompt
-        )
-
-    def select_model(self, model: str | None, temperature: float):
-        # Select and set the appropriate model based on the provided input
-        if model is None or model == "gpt-3.5-turbo":
-            print("Using model: gpt-3.5-turbo")
-            self.llm = ChatOpenAI(temperature=temperature)
-
-        if model == "text-davinci-003":
-            print("Using model: text-davinci-003")
-            self.llm = OpenAI(temperature=temperature)
-
-    def create_loader(self, index: str | None):
-        # Create a loader based on the provided directory (either local or S3)
-        if index is None:
-            raise RuntimeError(
-                """
-            Impossible to find a valid index. 
-            Either provide a valid path to a pickle file or a directory.               
-            """
-            )
-        self.loader = DirectoryLoader(index, recursive=True)
-
-    def load_or_create_index(self, index: str, vectorstore: VectorStore | None = None):
-        # Load an existing index from disk or create a new one if not available
-        if vectorstore is not None:
-            self.search_index = vectorstore.client.from_documents(
-                self.loader.load_and_split(),
-                OpenAIEmbeddings(),
-                connection_args={"host": vectorstore.host, "port": vectorstore.port},
-            )
-            return
-
-        # Is pickle
-        if index is not None and "pkl" in index or "pickle" in index:
-            print("Loading path from pickle file: ", index, "...")
-            with open(index, "rb") as f:
-                self.search_index = pickle.load(f)
-            return
-
-        # Is directory
-        if index is not None and os.path.isdir(index):
-            print("Creating index...")
-            self.search_index = FAISS.from_documents(
-                self.loader.load_and_split(), OpenAIEmbeddings()
-            )
-            return
-
-        raise RuntimeError(
-            """
-            Impossible to find a valid index. 
-            Either provide a valid path to a pickle file or a directory.               
-            """
-        )
-
-    def save_index(self, index_path: str):
-        # Save the index to the specified path
-        with open(index_path, "wb") as f:
-            pickle.dump(self.search_index, f)
-
-    def ask(self, question: str, k=1) -> str:
-        # Retrieve the answer to the given question and return it
-        input_documents = self.search_index.similarity_search(question, k=k)
-        answer = self.chain.run(input_documents=input_documents, question=question)
-        return answer
-
-
-SUPPORTED_TASKS = {
-    "qna-over-docs": {
-        "impl": Bot,
-        "default": {
-            "model": "gpt-3.5-turbo",
-            "temperature": 0,
-            "index": "./index",
-        },
-    }
-}
+from megabots.vectorstore import VectorStore, vectorstore
+from megabots.memory import Memory, memory
+from megabots.bot import Bot, bot
+from megabots.utils import create_api, create_interface
 
-SUPPORTED_MODELS = {}
 
-SUPPORTED_MEMORY = {
-    "conversation-buffer-window": {
-        "impl": ConversationBufferWindowMemory,
-        "default": {"memory_window": 3},
-    },
-    "conversation-buffer": {
-        "impl": ConversationBufferMemory,
-        "default": {},
-    },
-    "conversation-summary": {
-        "impl": ConversationSummaryMemory,
-        "default": {},
-    "conversation-summary-buffer": {
-        "impl": ConversationSummaryBufferMemory,
-        "default": {
-            "max_token_limit":40
-            }
-    },
-}
-
-
-def bot(
-    task: str | None = None,
-    model: str | None = None,
-    index: str | None = None,
-    prompt_template: str | None = None,
-    prompt_variables: list[str] | None = None,
-    memory: str | None = None,
-    memory_window: int = 3,
-    verbose: bool = False,
-    temperature: int = 0,
-    **kwargs,
-) -> Bot:
-    """Instanciate a bot based on the provided task. Each supported tasks has it's own default sane defaults.
-
-    Args:
-        task (str | None, optional): The given task. Can be one of the SUPPORTED_TASKS.
-        model (str | None, optional): Model to be used. Can be one of the SUPPORTED_MODELS.
-        index (str | None, optional): Data that the model will load and store index info.
-        Can be either a local file path, a pickle file, or a url of a vector database.
-        By default it will look for a local directory called "files" in the current working directory.
-        prompt_template (str | None, optional): Prompt template to be used. Specify variables with {var_name}.
-        prompt_variables (list[str] | None, optional): Prompt variables to be used in the prompt template.
-        verbose (bool, optional): Verbocity. Defaults to False.
-        temperature (int, optional): Temperature. Defaults to 0.
-
-    Raises:
-        RuntimeError: _description_
-        ValueError: _description_
-
-    Returns:
-        Bot: Bot instance
-    """
-
-    if task is None:
-        raise RuntimeError("Impossible to instantiate a bot without a task.")
-    if task not in SUPPORTED_TASKS:
-        raise ValueError(f"Task {task} is not supported.")
-
-    task_defaults = SUPPORTED_TASKS[task]["default"]
-    return SUPPORTED_TASKS[task]["impl"](
-        model=model or task_defaults["model"],
-        index=index or task_defaults["index"],
-        prompt_template=prompt_template,
-        prompt_variables=prompt_variables,
-        temperature=temperature,
-        verbose=verbose,
-        **kwargs,
-    )
-
-
-def create_api(bot: Bot):
-    app = FastAPI()
-
-    @app.get("/v1/ask/{question}")
-    async def ask(question: str):
-        answer = bot.ask(question)
-        return {"answer": answer}
-
-    return app
-
-
-def create_interface(bot_instance: Bot, examples: list[list[str]] = []):
-    with gr.Blocks() as interface:
-        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
-        msg = gr.Textbox(
-            show_label=False,
-            placeholder="Enter text and press enter, or upload an image",
-        ).style(container=False)
-        clear = gr.Button("Clear")
-
-        def user(user_message, history):
-            return "", history + [[user_message, None]]
-
-        def bot(history):
-            response = bot_instance.ask(history[-1][0])
-            history[-1][1] = response
-            return history
-
-        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot, chatbot, chatbot
-        )
-        clear.click(lambda: None, None, chatbot, queue=False)
+from dotenv import load_dotenv
 
-    return interface
+load_dotenv()

megabots/bot.py

@@ -0,0 +1,192 @@
+from typing import Any
+from langchain.llms import OpenAI
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.chains.qa_with_sources import load_qa_with_sources_chain
+from langchain.vectorstores.faiss import FAISS
+import pickle
+import os
+from langchain.prompts import PromptTemplate
+from langchain.chains.question_answering import load_qa_chain
+from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
+from langchain.document_loaders import DirectoryLoader
+from megabots.vectorstore import VectorStore
+from megabots.memory import Memory
+import megabots
+
+
+class Bot:
+    def __init__(
+        self,
+        model: str | None = None,
+        prompt_template: str | None = None,
+        prompt_variables: list[str] | None = None,
+        index: str | None = None,
+        sources: bool | None = False,
+        vectorstore: VectorStore | None = None,
+        memory: Memory | None = None,
+        verbose: bool = False,
+        temperature: int = 0,
+    ):
+        self.select_model(model, temperature)
+        self.create_loader(index)
+        self.load_or_create_index(index, vectorstore)
+        self.vectorstore = vectorstore
+        self.memory = memory
+        # Load the question-answering chain for the selected model
+        self.chain = self.create_chain(
+            prompt_template, prompt_variables, sources=sources, verbose=verbose
+        )
+
+    def create_chain(
+        self,
+        prompt_template: str | None = None,
+        prompt_variables: list[str] | None = None,
+        sources: bool | None = False,
+        verbose: bool = False,
+    ):
+        prompt = (
+            PromptTemplate(template=prompt_template, input_variables=prompt_variables)
+            if prompt_template is not None and prompt_variables is not None
+            else QA_PROMPT
+        )
+        # TODO: Changing the prompt here is not working. Leave it as is for now.
+        # Reference: https://github.com/hwchase17/langchain/issues/2858
+        if sources:
+            return load_qa_with_sources_chain(
+                self.llm, chain_type="stuff", verbose=verbose
+            )
+        return load_qa_chain(
+            self.llm, chain_type="stuff", verbose=verbose, prompt=prompt
+        )
+
+    def select_model(self, model: str | None, temperature: float):
+        # Select and set the appropriate model based on the provided input
+        if model is None or model == "gpt-3.5-turbo":
+            print("Using model: gpt-3.5-turbo")
+            self.llm = ChatOpenAI(temperature=temperature)
+
+        if model == "text-davinci-003":
+            print("Using model: text-davinci-003")
+            self.llm = OpenAI(temperature=temperature)
+
+    def create_loader(self, index: str | None):
+        # Create a loader based on the provided directory (either local or S3)
+        if index is None:
+            raise RuntimeError(
+                """
+            Impossible to find a valid index. 
+            Either provide a valid path to a pickle file or a directory.               
+            """
+            )
+        self.loader = DirectoryLoader(index, recursive=True)
+
+    def load_or_create_index(self, index: str, vectorstore: VectorStore | None = None):
+        # Load an existing index from disk or create a new one if not available
+        if vectorstore is not None:
+            self.search_index = vectorstore.client.from_documents(
+                self.loader.load_and_split(),
+                OpenAIEmbeddings(),
+                connection_args={"host": vectorstore.host, "port": vectorstore.port},
+            )
+            return
+
+        # Is pickle
+        if index is not None and "pkl" in index or "pickle" in index:
+            print("Loading path from pickle file: ", index, "...")
+            with open(index, "rb") as f:
+                self.search_index = pickle.load(f)
+            return
+
+        # Is directory
+        if index is not None and os.path.isdir(index):
+            print("Creating index...")
+            self.search_index = FAISS.from_documents(
+                self.loader.load_and_split(), OpenAIEmbeddings()
+            )
+            return
+
+        raise RuntimeError(
+            """
+            Impossible to find a valid index. 
+            Either provide a valid path to a pickle file or a directory.               
+            """
+        )
+
+    def save_index(self, index_path: str):
+        # Save the index to the specified path
+        with open(index_path, "wb") as f:
+            pickle.dump(self.search_index, f)
+
+    def ask(self, question: str, k=1) -> str:
+        # Retrieve the answer to the given question and return it
+        input_documents = self.search_index.similarity_search(question, k=k)
+        answer = self.chain.run(input_documents=input_documents, question=question)
+        return answer
+
+
+SUPPORTED_TASKS = {
+    "qna-over-docs": {
+        "impl": Bot,
+        "default": {
+            "model": "gpt-3.5-turbo",
+            "temperature": 0,
+            "index": "./index",
+        },
+    }
+}
+
+SUPPORTED_MODELS = {}
+
+
+def bot(
+    task: str | None = None,
+    model: str | None = None,
+    index: str | None = None,
+    prompt_template: str | None = None,
+    prompt_variables: list[str] | None = None,
+    memory: str | Memory | None = None,
+    vectorstore: str | VectorStore | None = None,
+    verbose: bool = False,
+    temperature: int = 0,
+) -> Bot:
+    """Instanciate a bot based on the provided task. Each supported tasks has it's own default sane defaults.
+
+    Args:
+        task (str | None, optional): The given task. Can be one of the SUPPORTED_TASKS.
+        model (str | None, optional): Model to be used. Can be one of the SUPPORTED_MODELS.
+        index (str | None, optional): Data that the model will load and store index info.
+        Can be either a local file path, a pickle file, or a url of a vector database.
+        By default it will look for a local directory called "files" in the current working directory.
+        prompt_template (str | None, optional): Prompt template to be used. Specify variables with {var_name}.
+        prompt_variables (list[str] | None, optional): Prompt variables to be used in the prompt template.
+        verbose (bool, optional): Verbocity. Defaults to False.
+        temperature (int, optional): Temperature. Defaults to 0.
+
+    Raises:
+        RuntimeError: _description_
+        ValueError: _description_
+
+    Returns:
+        Bot: Bot instance
+    """
+
+    if task is None:
+        raise RuntimeError("Impossible to instantiate a bot without a task.")
+    if task not in SUPPORTED_TASKS:
+        raise ValueError(f"Task {task} is not supported.")
+
+    task_defaults = SUPPORTED_TASKS[task]["default"]
+
+    return SUPPORTED_TASKS[task]["impl"](
+        model=model or task_defaults["model"],
+        index=index or task_defaults["index"],
+        prompt_template=prompt_template,
+        prompt_variables=prompt_variables,
+        temperature=temperature,
+        verbose=verbose,
+        vectorstore=megabots.vectorstore(vectorstore)
+        if isinstance(vectorstore, str)
+        else vectorstore,
+        memory=megabots.memory(memory) if isinstance(memory, str) else memory,
+    )

megabots/memory.py

@@ -0,0 +1,86 @@
+from langchain.memory import (
+    ConversationBufferMemory,
+    ConversationBufferWindowMemory,
+    ConversationSummaryMemory,
+    ConversationSummaryBufferMemory,
+)
+
+
+class ConversationBuffer:
+    def __init__(self):
+        self.memory = ConversationBufferMemory
+
+
+class ConversationBufferWindow:
+    def __init__(self, memory_window: int):
+        self.memory_window: int = memory_window
+        self.memory = ConversationBufferWindowMemory
+
+
+class ConversationSummary:
+    def __init__(self):
+        self.memory = ConversationSummaryMemory
+
+
+class ConversationSummaryBuffer:
+    def __init__(self, max_token_limit: int):
+        self.max_token_limit: int = max_token_limit
+        self.memory = ConversationSummaryBufferMemory
+
+
+SUPPORTED_MEMORY = {
+    "conversation-buffer": {
+        "impl": ConversationBuffer,
+        "default": {},
+    },
+    "conversation-buffer-window": {
+        "impl": ConversationBufferWindow,
+        "default": {"memory_window": 3},
+    },
+    "conversation-summary": {
+        "impl": ConversationSummary,
+        "default": {},
+    },
+    "conversation-summary-buffer": {
+        "impl": ConversationSummaryBuffer,
+        "default": {"max_token_limit": 40},
+    },
+}
+
+
+Memory = type(
+    "Memory",
+    (
+        ConversationBuffer,
+        ConversationBufferWindow,
+        ConversationSummary,
+        ConversationSummaryBuffer,
+    ),
+    {},
+)
+
+
+def memory(
+    name: str = "conversation-buffer-window",
+    memory_window: int | None = None,
+    max_token_limit: int | None = None,
+) -> Memory:
+    if name is None:
+        raise RuntimeError("Impossible to instantiate memory without a name.")
+
+    if name not in SUPPORTED_MEMORY:
+        raise ValueError(f"Memory {name} is not supported.")
+
+    cl = SUPPORTED_MEMORY[name]["impl"]
+
+    if name == "conversation-buffer-window":
+        if max_token_limit != None:
+            raise ValueError(f"max_token_limit cannot be set for {name} memory")
+        return cl(memory_window=memory_window)
+
+    if name == "conversation-summary-buffer":
+        if max_token_limit != None:
+            raise ValueError(f"memory_window cannot be set for {name} memory")
+        return cl(max_token_limit=max_token_limit)
+
+    return SUPPORTED_MEMORY[name]["impl"]()

megabots/utils.py

@@ -0,0 +1,39 @@
+import gradio as gr
+from fastapi import FastAPI
+from megabots.bot import Bot
+
+
+def create_api(bot: Bot):
+    app = FastAPI()
+
+    @app.get("/v1/ask/{question}")
+    async def ask(question: str):
+        answer = bot.ask(question)
+        return {"answer": answer}
+
+    return app
+
+
+def create_interface(bot_instance: Bot, examples: list[list[str]] = []):
+    with gr.Blocks() as interface:
+        chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
+        msg = gr.Textbox(
+            show_label=False,
+            placeholder="Enter text and press enter, or upload an image",
+        ).style(container=False)
+        clear = gr.Button("Clear")
+
+        def user(user_message, history):
+            return "", history + [[user_message, None]]
+
+        def bot(history):
+            response = bot_instance.ask(history[-1][0])
+            history[-1][1] = response
+            return history
+
+        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot, chatbot, chatbot
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    return interface

megabots/{vectorstores.py → vectorstore.py}

@@ -1,5 +1,5 @@
 from typing import Type, TypeVar
-from langchain.vectorstores import Milvus, Qdrant
+from langchain.vectorstores import Milvus
 from abc import ABC
 
 
@@ -26,7 +26,9 @@ SUPPORTED_VECTORSTORES = {
 }
 
 
-def vectorstore(name: str) -> VectorStore:
+def vectorstore(
+    name: str, host: str | None = None, port: int | None = None
+) -> VectorStore:
     """Return a vectorstore object."""
 
     if name is None:
@@ -36,6 +38,6 @@ def vectorstore(name: str) -> VectorStore:
         raise ValueError(f"Vectorstore {name} is not supported.")
 
     return SUPPORTED_VECTORSTORES[name]["impl"](
-        host=SUPPORTED_VECTORSTORES[name]["default"]["host"],
-        port=SUPPORTED_VECTORSTORES[name]["default"]["port"],
+        host=host or SUPPORTED_VECTORSTORES[name]["default"]["host"],
+        port=port or SUPPORTED_VECTORSTORES[name]["default"]["port"],
     )

tests/test_memory.py

@@ -0,0 +1,42 @@
+import pytest
+from megabots.memory import (
+    ConversationBufferWindow,
+    ConversationSummaryBuffer,
+    memory,
+    Memory,
+    SUPPORTED_MEMORY,
+)
+
+
+def test_memory_name_none():
+    with pytest.raises(RuntimeError):
+        memory(name=None)
+
+
+def test_memory_not_supported():
+    with pytest.raises(ValueError):
+        memory(name="unsupported_memory_type")
+
+
+def test_memory_conversation_buffer_window():
+    mem_obj = memory(name="conversation-buffer-window", memory_window=5)
+    assert isinstance(mem_obj, ConversationBufferWindow)
+    assert mem_obj.memory_window == 5
+    assert mem_obj.__class__ == SUPPORTED_MEMORY["conversation-buffer-window"]["impl"]
+
+
+def test_memory_conversation_buffer_window_invalid_max_token_limit():
+    with pytest.raises(ValueError):
+        memory(name="conversation-buffer-window", memory_window=5, max_token_limit=10)
+
+
+def test_memory_conversation_summary_buffer():
+    mem_obj = memory(name="conversation-summary-buffer", max_token_limit=10)
+    assert isinstance(mem_obj, ConversationSummaryBuffer)
+    assert mem_obj.max_token_limit == 10
+    assert mem_obj.__class__ == SUPPORTED_MEMORY["conversation-summary-buffer"]["impl"]
+
+
+def test_memory_conversation_summary_buffer_invalid_memory_window():
+    with pytest.raises(ValueError):
+        memory(name="conversation-summary-buffer", memory_window=5, max_token_limit=10)