/pentestgpt/modules/models.py

# models.py

from pentestgpt.prompts.prompt_class import PentestGPTPrompt
from pentestgpt.utils.chat_utils.agent import Agent


class ModelInteraction:
    def __init__(
        self,
        core,
        state,
        log_dir="logs",
        reasoning_model="gpt-4o",
        parsing_model="gpt-4-turbo",
        generation_model="gpt-4-turbo",
    ):
        self.core = core
        self.state = state
        self.log_dir = log_dir
        self.prompts = PentestGPTPrompt()
        self.task_description = None
        self.previous_user_input = None

        self.initialization(reasoning_model, parsing_model, generation_model)

    def initialization(self, reasoning_model, parsing_model, generation_model):
        # Initialization messages are handled by core
        self.reasoning_agent = Agent(
            "reasoning", reasoning_model, self.prompts.reasoning_session_init
        )
        self.parsing_agent = Agent(
            "parsing", parsing_model, self.prompts.input_parsing_init
        )
        self.generation_agent = Agent(
            "generation", generation_model, self.prompts.generation_session_init
        )

    def search_information(self, agent, keyword):
        # require a specific agent to perform information search
        return agent.search(keyword)

    def process(self, user_input: str, task_description: str) -> list:
        response_messages = []

        # Update overall task description to the reasoning agent
        if task_description != self.task_description:
            self.task_description = task_description
            self.state.set_variable("task_description", task_description)

            # Initial reasoning agent chat with spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.reasoning_agent.name}...",
                spinner="dots",
            ):
                _ = self.reasoning_agent.chat(
                    self.prompts.task_description + task_description
                )

            if self.previous_user_input is not None:
                response_messages.append(
                    {
                        "type": "info",
                        "message": "Task description updated successfully. Please continue to update your input.",
                    }
                )
            else:
                response_messages.append(
                    {
                        "type": "info",
                        "message": "Task description updated successfully. Now generating the detailed task.",
                    }
                )

                # Task selection with spinner
                with self.core.console.status(
                    f"[bold green]Processing with {self.reasoning_agent.name}...",
                    spinner="dots",
                ):
                    selected_task = self.reasoning_agent.chat(
                        self.prompts.process_results_task_selection
                    )

                # Final task generation with spinner
                with self.core.console.status(
                    f"[bold green]Processing with {self.generation_agent.name}...",
                    spinner="dots",
                ):
                    final_task = self.generation_agent.chat(
                        self.prompts.todo_to_command + selected_task
                    )

                response_messages.extend(
                    [
                        {"type": "heading", "message": "Initial Task Tree"},
                        {"type": "heading", "message": "Final Task"},
                        {"type": "text", "message": final_task},
                    ]
                )

            return response_messages

        # Process the user input
        if user_input == self.previous_user_input:
            response_messages.append(
                {
                    "type": "warning",
                    "message": "It seems you have already provided this input in the previous iteration. Please provide a new input.",
                }
            )
        else:
            # 1. Parse the user input with spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.parsing_agent.name}...",
                spinner="dots",
            ):
                parsing_result = self.parsing_agent.query(
                    self.prompts.input_parsing_command + user_input
                )

            # 2. Update PTT with spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.reasoning_agent.name}...",
                spinner="dots",
            ):
                updated_tree = self.reasoning_agent.chat(
                    self.prompts.process_results + parsing_result
                )

            # 3. Select the next task
            # Obtain keywords with spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.reasoning_agent.name} to obtain keywords...",
                spinner="dots",
            ):
                keywords = self.reasoning_agent.query(
                    self.prompts.generation_obtain_keywords
                )

            # Perform online search with spinner
            with self.core.console.status(
                "[bold green]Performing online search...", spinner="dots"
            ):
                additional_context = self.search_information(
                    self.reasoning_agent, keywords
                )

            # Reasoning agent chat with additional context and spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.reasoning_agent.name}...",
                spinner="dots",
            ):
                selected_task = self.reasoning_agent.chat(
                    self.prompts.process_results_task_selection,
                    additional_context=additional_context,
                )

            # 4. Final task generation with spinner
            with self.core.console.status(
                f"[bold green]Processing with {self.generation_agent.name}...",
                spinner="dots",
            ):
                final_task = self.generation_agent.chat(
                    self.prompts.todo_to_command + updated_tree + "\n" + selected_task
                )

            self.previous_user_input = user_input

            response_messages.extend(
                [
                    {"type": "heading", "message": "Current Task Tree"},
                    {"type": "text", "message": updated_tree},
                    {"type": "heading", "message": "Final Task"},
                    {"type": "text", "message": final_task},
                ]
            )

        return response_messages