Upload bot.py

bot.py

@@ -0,0 +1,200 @@
+import logging
+from pathlib import Path
+from typing import List, Tuple
+
+import fire
+
+logger = logging.getLogger(__name__)
+
+# === Bot Loaders ===
+
+
+def load_bot(
+    env_file_path: str = ".env",
+    logging_config_path: str = "logging.yaml",
+    model_cache_dir: str = "/model_cache",
+    embedding_model_device: str = "cuda:0",
+    debug: bool = False,
+):
+    """
+    Load the financial assistant bot in production or development mode based on the `debug` flag
+
+    In DEV mode the embedding model runs on CPU and the fine-tuned LLM is mocked.
+    Otherwise, the embedding model runs on GPU and the fine-tuned LLM is used.
+
+    Args:
+        env_file_path (str): Path to the environment file.
+        logging_config_path (str): Path to the logging configuration file.
+        model_cache_dir (str): Path to the directory where the model cache is stored.
+        embedding_model_device (str): Device to use for the embedding model.
+        debug (bool): Flag to indicate whether to run the bot in debug mode or not.
+
+    Returns:
+        FinancialBot: An instance of the FinancialBot class.
+    """
+
+    from financial_bot import initialize
+
+    # Be sure to initialize the environment variables before importing any other modules.
+    initialize(logging_config_path=logging_config_path, env_file_path=env_file_path)
+
+    from financial_bot import utils
+    from financial_bot.langchain_bot import FinancialBot
+
+    logger.info("#" * 100)
+    utils.log_available_gpu_memory()
+    utils.log_available_ram()
+    logger.info("#" * 100)
+
+    bot = FinancialBot(
+        model_cache_dir=Path(model_cache_dir) if model_cache_dir else None,
+        embedding_model_device=embedding_model_device,
+        debug=debug,
+    )
+
+    return bot
+
+
+def load_bot_dev(
+    env_file_path: str = ".env",
+    logging_config_path: str = "logging.yaml",
+    model_cache_dir: str = "./model_cache",
+):
+    """
+    Load the Financial Assistant Bot in dev mode: the embedding model runs on CPU and the LLM is mocked.
+
+    Args:
+        env_file_path (str): Path to the environment file.
+        logging_config_path (str): Path to the logging configuration file.
+        model_cache_dir (str): Path to the directory where the model cache is stored.
+
+    Returns:
+        The loaded Financial Assistant Bot in dev mode.
+    """
+
+    return load_bot(
+        env_file_path=env_file_path,
+        logging_config_path=logging_config_path,
+        model_cache_dir=model_cache_dir,
+        embedding_model_device="cpu",
+        debug=True,
+    )
+
+
+# === Bot Runners ===
+
+
+@financial_bot.rest_api(keep_warm_seconds=300, loader=load_bot)
+def run(**inputs):
+    """
+    Run the bot under the Beam RESTful API endpoint.
+
+     Args:
+        inputs (dict): A dictionary containing the following keys:
+            - context: The bot instance.
+            - about_me (str): Information about the user.
+            - question (str): The user's question.
+            - history (list): A list of previous conversations (optional).
+
+    Returns:
+        str: The bot's response to the user's question.
+    """
+
+    response = _run(**inputs)
+
+    return response
+
+
+@financial_bot_dev.rest_api(keep_warm_seconds=300, loader=load_bot_dev)
+def run_dev(**inputs):
+    """
+    Run the bot under the Beam RESTful API endpoint [Dev Mode].
+
+     Args:
+        inputs (dict): A dictionary containing the following keys:
+            - context: The bot instance.
+            - about_me (str): Information about the user.
+            - question (str): The user's question.
+            - history (list): A list of previous conversations (optional).
+
+    Returns:
+        str: The bot's response to the user's question.
+    """
+
+    response = _run(**inputs)
+
+    return response
+
+
+def run_local(
+    about_me: str,
+    question: str,
+    history: List[Tuple[str, str]] = None,
+    debug: bool = False,
+):
+    """
+    Run the bot locally in production or dev mode.
+
+    Args:
+        about_me (str): A string containing information about the user.
+        question (str): A string containing the user's question.
+        history (List[Tuple[str, str]], optional): A list of tuples containing the user's previous questions
+            and the bot's responses. Defaults to None.
+        debug (bool, optional): A boolean indicating whether to run the bot in debug mode. Defaults to False.
+
+    Returns:
+        str: A string containing the bot's response to the user's question.
+    """
+
+    if debug is True:
+        bot = load_bot_dev(model_cache_dir=None)
+    else:
+        bot = load_bot(model_cache_dir=None)
+
+    inputs = {
+        "about_me": about_me,
+        "question": question,
+        "history": history,
+        "context": bot,
+    }
+
+    response = _run(**inputs)
+
+    return response
+
+
+def _run(**inputs):
+    """
+    Central function that calls the bot and returns the response.
+
+    Args:
+        inputs (dict): A dictionary containing the following keys:
+            - context: The bot instance.
+            - about_me (str): Information about the user.
+            - question (str): The user's question.
+            - history (list): A list of previous conversations (optional).
+
+    Returns:
+        str: The bot's response to the user's question.
+    """
+
+    from financial_bot import utils
+
+    logger.info("#" * 100)
+    utils.log_available_gpu_memory()
+    utils.log_available_ram()
+    logger.info("#" * 100)
+
+    bot = inputs["context"]
+    input_payload = {
+        "about_me": inputs["about_me"],
+        "question": inputs["question"],
+        "to_load_history": inputs["history"] if "history" in inputs else [],
+    }
+    response = bot.answer(**input_payload)
+
+    return response
+
+
+if __name__ == "__main__":
+    fire.Fire(run_local)