Refactor base_telegram_inference_bot.py: Improve error handling, encapsulate processing status, and define abstract methods.

base_telegram_inference_bot.py

@@ -2,6 +2,7 @@ import importlib
 import os
 import json
 import inspect
+import logging
 from abc import ABC, abstractmethod
 from tools.base_tool import BaseTool
 
@@ -11,32 +12,44 @@ class BaseTelegramInferenceBot(ABC):
         self.processing_status = {}
         self.system_prompt = self.load_system_prompt()
         self.tools, self.functions = self.load_functions()
-        print(f'System Prompt: {os.environ.get("SYSTEM_PROMPT_PATH")}')
-        print(f'Github Repository: {os.environ.get("GITHUB_REPOSITORY")}')
+        logging.info(f'System Prompt: {os.environ.get("SYSTEM_PROMPT_PATH")}')
+        logging.info(f'Github Repository: {os.environ.get("GITHUB_REPOSITORY")}')
 
-    @staticmethod
-    def load_system_prompt():
+    def load_system_prompt(self):
         system_prompt_path = os.getenv("SYSTEM_PROMPT_PATH")
         if system_prompt_path and os.path.isfile(system_prompt_path):
-            with open(system_prompt_path, "r", encoding="utf-8") as file:
-                return file.read().strip()
+            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:
-            raise FileNotFoundError("SYSTEM_PROMPT_PATH is not set or file does not exist.")
+            logging.warning("SYSTEM_PROMPT_PATH is not set or file does not exist. Using default system prompt.")
+            return "You are a helpful AI assistant."
 
-    @staticmethod
-    def load_functions():
+    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]}'
-                module = importlib.import_module(module_name)
-                for name, obj in inspect.getmembers(module):
-                    if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool:
-                        tools.append(obj())
+                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}")
         
-        # Collect all function definitions
-        functions = []
         for tool in tools:
             functions.extend(tool.get_functions())
         return tools, functions
@@ -49,47 +62,64 @@ class BaseTelegramInferenceBot(ABC):
     async def handle_message(self, user_id, user_message):
         pass
 
-    def clear_conversation(self, user_id):
+    def clear_conversation_history(self, user_id):
         if user_id in self.conversation_history:
             del self.conversation_history[user_id]
         
+        # Assuming tool.clear() is for global state or doesn't need 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
-        function_args = json.loads(function_call_arguments if function_call_arguments is not None else "{}")
+        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:
-                    return tool.execute(function_name, **function_args)
+                    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."
     
-    @abstractmethod
     def get_system_prompt_description(self) -> str:
         """Returns a description of the system prompt being used."""
-        pass
+        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 status(self) -> str: # Changed from abstract to concrete
+    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()
-        # Consider potential async calls if get_... methods were async
-        # For now, assuming they are synchronous as per design
         return f"{prompt_desc}\n{llm_desc}"
 
     @abstractmethod
     async def start(self):
         pass
 
-    @abstractmethod
-    async def clear(self, user_id):
-        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