cyclop/base_telegram_inference_bot.py

import importlib
import os
import json
import inspect
import logging
from abc import ABC, abstractmethod
from tools.base_tool import BaseTool

class BaseTelegramInferenceBot(ABC):
    def __init__(self):
        self.conversation_history = {}
        self.processing_status = {}
        self.system_prompt = self.load_system_prompt()
        self.tools, self.functions = self.load_functions()
        logging.info(f'System Prompt: {os.environ.get("SYSTEM_PROMPT_PATH")}')
        logging.info(f'Github Repository: {os.environ.get("GITHUB_REPOSITORY")}')

    def load_system_prompt(self):
        system_prompt_path = os.getenv("SYSTEM_PROMPT_PATH")
        if system_prompt_path and os.path.isfile(system_prompt_path):
            try:
                with open(system_prompt_path, "r", encoding="utf-8") as file:
                    return file.read().strip()
            except IOError as e:
                logging.warning(f"Could not read system prompt file {system_prompt_path}: {e}")
                return "You are a helpful AI assistant."
        else:
            logging.warning("SYSTEM_PROMPT_PATH is not set or file does not exist. Using default system prompt.")
            return "You are a helpful AI assistant."

    def load_functions(self):
        tools = []
        functions = []
        tools_dir = os.path.join(os.path.dirname(__file__), 'tools')
        if not os.path.exists(tools_dir):
            logging.warning(f"Tools directory not found: {tools_dir}")
            return [], []

        for filename in os.listdir(tools_dir):
            if filename.endswith('.py') and filename != '__init__.py' and filename != 'base_tool.py':
                module_name = f'tools.{filename[:-3]}'
                try:
                    module = importlib.import_module(module_name)
                    for name, obj in inspect.getmembers(module):
                        if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool:
                            try:
                                tools.append(obj())
                            except Exception as e:
                                logging.error(f"Error instantiating tool {name} from {filename}: {e}")
                except Exception as e:
                    logging.error(f"Error importing module {module_name}: {e}")
        
        for tool in tools:
            functions.extend(tool.get_functions())
        return tools, functions

    @abstractmethod
    def get_chat_response(self, messages):
        pass

    @abstractmethod
    async def handle_message(self, user_id, user_message):
        pass

    def clear_conversation_history(self, user_id):
        if user_id in self.conversation_history:
            del self.conversation_history[user_id]
        
        for tool in self.tools:
            tool.clear()

    def set_processing_status(self, user_id: int, message_id: int):
        self.processing_status[user_id] = {"processing": True, "message_id": message_id}

    def clear_processing_status(self, user_id: int):
        if user_id in self.processing_status:
            del self.processing_status[user_id]

    def call_tool(self, function_call_name, function_call_arguments):
        function_name = function_call_name
        try:
            function_args = json.loads(function_call_arguments if function_call_arguments is not None else "{}")
        except json.JSONDecodeError as e:
            logging.error(f"Error decoding function call arguments for {function_call_name}: {e}. Arguments: {function_call_arguments}")
            return f"Error: Malformed arguments for tool call: {e}"

        for tool in self.tools:
            for function in tool.get_functions():
                if function["function"]["name"] == function_name:
                    try:
                        return tool.execute(function_name, **function_args)
                    except Exception as e:
                        logging.error(f"Error executing tool {function_name} with args {function_args}: {e}")
                        return f"Error executing tool {function_name}: {e}"
        logging.warning(f"Tool function {function_name} not found.")
        return f"Error: Tool function {function_name} not found."
    
    def get_system_prompt_description(self) -> str:
        """Returns a description of the system prompt being used."""
        return f"System Prompt: {'Custom' if os.getenv('SYSTEM_PROMPT_PATH') else 'Default'}"

    @abstractmethod
    def get_llm_description(self) -> str:
        """Returns a description of the LLM being used."""
        pass

    async def get_bot_status(self) -> str:
        """Provides a status message including prompt and LLM information."""
        prompt_desc = self.get_system_prompt_description()
        llm_desc = self.get_llm_description()
        return f"{prompt_desc}\n{llm_desc}"

    @abstractmethod
    async def start(self):
        pass

    @abstractmethod
    async def abort_processing(self, user_id):
        pass
    
    @abstractmethod
    async def switch_model(self):
        """Switches the underlying model if supported by the bot."""
        pass