import dataclasses import importlib import os import sys module_mapping = { "gpt-4o": { "config_name": "GPT4O", "module_name": "chatgpt_api", "class_name": "ChatGPTAPI", }, "o3": { # Added GPT-o3 "config_name": "GPT4O3", "module_name": "chatgpt_api", "class_name": "ChatGPTAPI", }, "o4-mini": { # Added GPT-o4-mini "config_name": "GPT4O4Mini", "module_name": "chatgpt_api", "class_name": "ChatGPTAPI", }, "gpt4all": { "config_name": "GPT4ALLConfigClass", "module_name": "gpt4all_api", "class_name": "GPT4ALLAPI", }, "gemini-2.5-flash": { # Added Gemini 2.5 "config_name": "Gemini25ConfigClass", "module_name": "gemini_api", "class_name": "GeminiAPI", }, "gemini-2.5-pro": { # Added Gemini 2.5 "config_name": "Gemini25ProConfigClass", "module_name": "gemini_api", "class_name": "GeminiAPI", }, "deepseek-r1": { # Added Deepseek R1 "config_name": "DeepseekR1ConfigClass", "module_name": "deepseek_api", "class_name": "DeepseekAPI", }, "deepseek-v3": { # Added Deepseek v3 "config_name": "DeepseekV3ConfigClass", "module_name": "deepseek_api", "class_name": "DeepseekAPI", }, "ollama": { # Added Ollama support "config_name": "OllamaConfigClass", "module_name": "ollama_api", "class_name": "OllamaAPI", }, } @dataclasses.dataclass class GPT4O: model: str = "gpt-4o" api_base: str = os.getenv("OPENAI_BASEURL", "https://api.openai.com/v1") # set up the openai key openai_key = os.getenv("OPENAI_API_KEY", None) if openai_key is None: print( "Your OPENAI_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 10 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class GPT4O3: # Added GPT-o3 configuration model: str = "o3" api_base: str = os.getenv("OPENAI_BASEURL", "https://api.openai.com/v1") # set up the openai key openai_key = os.getenv("OPENAI_API_KEY", None) if openai_key is None: print( "Your OPENAI_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 10 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class GPT4O4Mini: # Added GPT-o4-mini configuration model: str = "o4-mini" api_base: str = os.getenv("OPENAI_BASEURL", "https://api.openai.com/v1") # set up the openai key openai_key = os.getenv("OPENAI_API_KEY", None) if openai_key is None: print( "Your OPENAI_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 10 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class GPT4ALLConfigClass: model: str = "mistral-7b-openorca.Q4_0.gguf" @dataclasses.dataclass class TitanConfigClass: model: str = "amazon.titan-tg1-large" @dataclasses.dataclass class Gemini25ConfigClass: # Added Gemini 2.5 flash configuration model: str = "gemini-2.5-flash-preview-04-17" gemini_key = os.getenv( "GEMINI_API_KEY", None ) # Assuming environment variable for API key if gemini_key is None: print( "Your GEMINI_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 20 is_debugging: bool = False log_dir: str = None class Gemini25ProConfigClass: # Added Gemini 2.5 Pro configuration model: str = "gemini-2.5-pro-preview-03-25" gemini_key = os.getenv( "GEMINI_API_KEY", None ) # Assuming environment variable for API key if gemini_key is None: print( "Your GEMINI_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 20 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class DeepseekR1ConfigClass: # Added Deepseek configuration model: str = "deepseek-reasoner" api_base: str = os.getenv("DEEPSEEK_BASEURL", "https://api.deepseek.com") deepseek_key = os.getenv("DEEPSEEK_API_KEY", None) if deepseek_key is None: print( "Your DEEPSEEK_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 20 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class DeepseekV3ConfigClass: # Added Deepseek configuration model: str = "deepseek-chat" api_base: str = os.getenv("DEEPSEEK_BASEURL", "https://api.deepseek.com") deepseek_key = os.getenv("DEEPSEEK_API_KEY", None) if deepseek_key is None: print( "Your DEEPSEEK_API_KEY is not set. Please set it in the environment variable." ) error_wait_time: float = 20 is_debugging: bool = False log_dir: str = None @dataclasses.dataclass class OllamaConfigClass: # Added Ollama configuration model: str = os.getenv("OLLAMA_MODEL", "llama3.1:latest") # Default to more common model api_base: str = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") # Default Ollama URL error_wait_time: float = 10 is_debugging: bool = False log_dir: str = None def dynamic_import(module_name, log_dir, use_langfuse_logging=False) -> object: if module_name in module_mapping: module_config_name = module_mapping[module_name]["config_name"] module_import_name = module_mapping[module_name]["module_name"] class_name = module_mapping[module_name]["class_name"] module_config = getattr(sys.modules[__name__], module_config_name) module_config.log_dir = log_dir # Special handling for Ollama to use the model from environment variable if module_name == "ollama" and "OLLAMA_MODEL" in os.environ: module_config.model = os.environ["OLLAMA_MODEL"] # import the module LLM_module = importlib.import_module( "pentestgpt.utils.APIs." + module_import_name ) LLM_class = getattr(LLM_module, class_name) # initialize the class LLM_class_initialized = LLM_class( module_config, use_langfuse_logging=use_langfuse_logging ) return LLM_class_initialized else: print( "Module not found: " + module_name + ". Falling back to use the default gpt-4o" ) # fall back to gpt-3.5-turbo-16k LLM_class_initialized = dynamic_import("gpt-4o", log_dir) return LLM_class_initialized if __name__ == "__main__": # Quick test for each model import time # Define test models based on the revised list test_models = [ # "gpt-4o", # "o3", # "o4-mini", # "gemini-2.5-flash", # "gemini-2.5-pro", "deepseek-r1", "deepseek-v3", ] # Test message test_message = "Explain briefly what makes you unique as an AI model." print("=== Starting Model Tests ===") for model_name in test_models: try: print(f"\n\nTesting model: {model_name}") print("=" * 50) # Initialize the model llm = dynamic_import(model_name, "logs") # Print model information print(f"Model initialized: {llm.__class__.__name__}") # Send a message and get response print(f"\nSending test message: '{test_message}'") print("-" * 50) # Measure response time start_time = time.time() response = llm.send_new_message(test_message) end_time = time.time() # Print response and timing print(f"Response received in {end_time - start_time:.2f} seconds:") print("-" * 50) print(response) print("-" * 50) # Add a small delay between API calls time.sleep(2) except Exception as e: print(f"Error testing {model_name}: {str(e)}") print("\n=== Model Tests Complete ===")