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, system_prompt_content: str | None = None, system_prompt_path: str | None = None): # MODIFIED
        self.conversation_history = {}
        self.processing_status = {}
        # MODIFIED to pass arguments
        self.system_prompt = self.load_system_prompt(
            direct_content=system_prompt_content,
            file_path=system_prompt_path
        )
        self.tools, self.functions = self.load_functions()
        # Logging the actual source of the system prompt might be more complex now,
        # but we can log the final prompt or indicate if it's custom/default.
        # We'll also log the source of the prompt inside load_system_prompt.
        logging.info(f'System Prompt (effective): {"Custom" if self.system_prompt != "You are a helpful AI assistant." else "Default"}')
        logging.info(f'Github Repository: {os.environ.get("GITHUB_REPOSITORY")}')

    def load_system_prompt(self, direct_content: str | None = None, file_path: str | None = None) -> str: # MODIFIED
        default_prompt = "You are a helpful AI assistant."

        if direct_content:
            logging.info("Using direct content for system prompt.")
            return direct_content.strip()

        prompt_path_to_try = file_path or os.getenv("SYSTEM_PROMPT_PATH")

        if prompt_path_to_try:
            if os.path.isfile(prompt_path_to_try):
                try:
                    with open(prompt_path_to_try, "r", encoding="utf-8") as file:
                        content = file.read().strip()
                        logging.info(f"Successfully loaded system prompt from {prompt_path_to_try}.")
                        return content
                except IOError as e:
                    logging.warning(f"Could not read system prompt file {prompt_path_to_try}: {e}. Using default.")
                    return default_prompt
            else:
                # This condition now also covers if 'file_path' argument was given but invalid
                logging.warning(f"System prompt file {prompt_path_to_try} not found. Using default system prompt.")
                return default_prompt
        else:
            logging.info("No system prompt path provided (argument or ENV) or direct content. Using default system prompt.")
            return default_prompt

    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()) # This instantiation might be an issue for tools needing config
                            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
        function_args = None
        if isinstance(function_call_arguments, dict):
            function_args = function_call_arguments
        elif isinstance(function_call_arguments, str):
            try:
                function_args = json.loads(function_call_arguments)
            except json.JSONDecodeError as e:
                logging.error(f"Error decoding function call arguments (string) for {function_call_name}: {e}. Arguments: {function_call_arguments}")
                return f"Error: Malformed arguments for tool call: {e}"
        else: 
            if function_call_arguments is None:
                function_args = {} 
            else:
                logging.error(f"Unexpected type for function_call_arguments for {function_call_name}: {type(function_call_arguments)}. Arguments: {function_call_arguments}")
                return f"Error: Invalid argument type for tool call: {type(function_call_arguments)}"

        for tool in self.tools:
            for function in tool.get_functions():
                if function["function"]["name"] == function_name:
                    try:
                        if not isinstance(function_args, dict):
                             logging.error(f"Internal error: function_args not a dict for {function_name} before execution. Args: {function_args}")
                             return f"Internal error preparing arguments for tool {function_name}."
                        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:
        # This method could be updated to be more specific about the prompt source if needed.
        # For now, it still reflects custom vs default based on the original ENV var logic's spirit.
        # A more accurate reflection would require storing how the prompt was loaded.
        # For simplicity, let's assume if it's not the default, it's "Custom".
        if self.system_prompt != "You are a helpful AI assistant.":
            return "System Prompt: Custom"
        # Check original ENV var for backward compatibility in description only
        elif os.getenv('SYSTEM_PROMPT_PATH'):
             return "System Prompt: Custom (via ENV)"
        return "System Prompt: Default"


    @abstractmethod
    def get_llm_description(self) -> str:
        pass

    async def get_bot_status(self) -> str:
        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):
        pass