import os import sys from abc import ABC from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate from Lib.configs import DATA_DIR from Lib.log import logger class BaseAPI(ABC): def __init__(self): self.logger = logger class _TemplateWrapper: """ A template wrapper class hidden internally that does only one thing: provide .format().""" def __init__(self, content: str): self._content = content def format(self, **kwargs) -> str: """ Implement the .format() method you want. """ return self._content.format(**kwargs) @staticmethod def _get_main_script_name(): """ Get the filename of the main execution script (without the extension). sys.argv[0] always points to the script that was originally started, regardless of which module the current code is running in. """ try: # 1. Get the full path of the main execution script script_path = sys.argv[0] # 2. Extract the file name from the full path script_filename = os.path.basename(script_path) # 3. Separate the file name and extension script_name, _ = os.path.splitext(script_filename) return script_name except IndexError as e: raise RuntimeError("Unable to get the name of the main execution script, sys.argv[0] does not exist.") from e except Exception as e: raise RuntimeError(f"An error occurred while getting the name of the main execution script: {e}") from e @property def module_name(self): """Get the module loading path""" module_name = self.__module__.split(".")[-1] if module_name == "__main__": return self._get_main_script_name() else: return module_name def _get_md_file_path(self, filename: str, lang=None) -> str: """ Get the file path based on the workbook name. """ if os.path.isfile(filename): # "/root/asf/ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" template_path = filename else: if filename.endswith('.md'): # "senior_phishing_expert.md" fname = filename else: if lang is not None: fname = f"{filename}_{lang}.md" # "senior_phishing_expert_en" else: fname = f"{filename}.md" # "senior_phishing_expert" if os.path.isfile(os.path.join(DATA_DIR, fname)): # "ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" template_path = os.path.join(DATA_DIR, fname) else: template_path = os.path.join(DATA_DIR, self.module_name, fname) return template_path def _get_file_path(self, filename: str): """ Get the file path based on the workbook name. """ if os.path.isfile(filename): # "/root/asf/ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" return filename else: if os.path.join(DATA_DIR, self.module_name, filename): # "ES-Rule-21-Phishing_user_report_mail/senior_phishing_expert.md" template_path = os.path.join(DATA_DIR, self.module_name, filename) return template_path else: raise Exception("File not exist") def load_markdown_template(self, filename: str) -> _TemplateWrapper: """ Read the content according to the workbook name and return an object that supports .format(). """ template_path = self._get_md_file_path(filename) try: with open(template_path, 'r', encoding='utf-8') as f: content = f.read() # Return an instance of an inner nested class return self._TemplateWrapper(content) except Exception as e: logger.warning(f"Failed to load prompt template {template_path}: {str(e)}") raise e def load_system_prompt_template(self, filename, lang=None): """Load system prompt template""" template_path = self._get_md_file_path(filename, lang=lang) try: with open(template_path, 'r', encoding='utf-8') as f: system_prompt_template: SystemMessagePromptTemplate = SystemMessagePromptTemplate.from_template(f.read()) logger.debug(f"Loaded system prompt template from: {template_path}") return system_prompt_template except Exception as e: logger.warning(f"Failed to load prompt template {template_path}: {str(e)}") raise e def load_human_prompt_template(self, filename, lang=None): template_path = self._get_md_file_path(filename, lang=lang) try: with open(template_path, 'r', encoding='utf-8') as f: human_prompt_template: HumanMessagePromptTemplate = HumanMessagePromptTemplate.from_template(f.read()) logger.debug(f"Loaded human prompt template from: {template_path}") return human_prompt_template except Exception as e: logger.warning(f"Failed to load prompt template {template_path}: {str(e)}") raise e def run(self): raise NotImplementedError