agentic-soc-platform/PLAYBOOKS/Threat_Hunting_Agent.py

import json
import operator
from typing import Annotated, Dict, List

from langchain_core.messages import AnyMessage, ToolMessage, AIMessage, HumanMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt import ToolNode
from langgraph.types import Send
from pydantic import BaseModel, Field, ConfigDict

from AGENTS.agent_siem import AgentSIEM
from AGENTS.agent_threat_intelligence import AgentThreatIntelligence
from Lib.api import get_current_time_str
from Lib.baseplaybook import LanggraphPlaybook
from PLUGINS.LLM.llmapi import LLMAPI
from PLUGINS.SIRP.sirpapi import Case
from PLUGINS.SIRP.sirpbasemodel import AI_PROFILE_INVESTIGATION
from PLUGINS.SIRP.sirpcoremodel import CaseModel
from PLUGINS.SIRP.sirpextramodel import PlaybookJobStatus, PlaybookModel

MAX_ITERATIONS = 3
MAX_ITERATIONS_OF_FUNCTIONS_CALL = 2
PROMPT_LANG = None
tools = [
    AgentSIEM.siem_search_by_natural_language,
    AgentThreatIntelligence.threat_intelligence_lookup
]


class PlanningRecord(BaseModel):
    """Structured record for storing single-round planning"""
    iteration: int = Field(description="The current round")
    rationale: str = Field(description="The reasoning for the plan")
    plan: List[str] = Field(description="The specific list of tasks generated")

    def to_markdown(self) -> str:
        tasks_str = ", ".join(self.plan)
        return (f"#### Round {self.iteration}\n"
                f"**Reasoning:** {self.rationale}\n"
                f"**Tasks Executed:** {tasks_str}\n")


class AnalystOutput(BaseModel):
    answer: str = Field(description="The final, concise answer to the investigation question")
    reasoning: str | List[str] = Field(description="The detailed reasoning process and key evidence supporting the final conclusion.")


class HuntingPlan(BaseModel):
    current_plan: List[str] = Field(
        description="A list of specific questions to be investigated in parallel next. Returns an empty list if there are no more questions.")
    rationale: str = Field(description="The reason for making this plan")


class Finding(BaseModel):
    question: str = Field(description="The question to be investigated")
    answer: str = Field(description="The answer obtained from the investigation")
    reasoning: str | List[str] = Field(description="The reasoning process of the investigation")
    tool_calls: List = Field(default_factory=list, description="Tool call records")

    def to_markdown(self) -> str:
        return (f"\n"
                f"**Question:** {self.question}\n"
                f"**Reasoning:** {self.reasoning}\n"
                f"**Answer:** {self.answer}\n"
                f"\n\n\n"
                )


class AnalystState(BaseModel):
    """
    Subgraph state: responsible for the execution of a single investigation task.
    Inherits from Pydantic BaseModel, supports default values and data validation.
    """
    model_config = ConfigDict(
        arbitrary_types_allowed=True
    )
    # Message history
    messages: Annotated[
        List[AnyMessage],
        add_messages
    ] = Field(
        default_factory=list,
        description="Used to store the message passing history between all nodes."
    )

    question: str = Field(
        description="The specific question to be investigated, usually the starting point of user input."
    )

    case: CaseModel = Field(
        description="External context or additional data provided for this investigation task."
    )

    answer: str = Field(
        default="",
        description="The final conclusion or summary of the investigation task."
    )

    reasoning: str | List[str] = Field(
        default="",
        description="Detailed reasoning steps and evidence for the conclusion."
    )

    tool_calls: Annotated[
        List[Dict],
        operator.add
    ] = Field(
        default_factory=list,
        description="A log of tool calls and their results."
    )
    loop_count: int = Field(
        default=0,
        description="Count of function call iterations."
    )


class MainState(BaseModel):
    """
    [Main graph state]
    Responsible for global planning and summarization.
    """
    model_config = ConfigDict(
        arbitrary_types_allowed=True
    )

    case: CaseModel = Field(
        default_factory=dict,
        description="Original case or global context data."
    )

    user_intent: str = Field(
        default="",
        description="The initial request or core intent proposed by the user."
    )

    hunting_objective: str = Field(
        default="",
        description="The overall goal that the entire graph needs to achieve, established based on the user's intent."
    )

    findings: Annotated[
        List[Finding],
        operator.add
    ] = Field(
        default_factory=list,
        description="A list of all results or findings collected from subgraphs or subtasks."
    )

    current_plan: List[str] = Field(
        default_factory=list,
        description="A list of planned tasks to be executed in the current iteration or batch."
    )

    iteration_count: int = Field(
        default=0,
        description="A count of the number of main graph loops or iterations."
    )

    planning_history: Annotated[
        List[PlanningRecord],
        operator.add
    ] = Field(
        default_factory=list,
        description="Structured reasoning history record."
    )

    # Final output
    report: str = Field(
        default="",
        description="The final report summarized and organized based on all Findings."
    )


class Playbook(LanggraphPlaybook):
    NAME = "Threat Hunting Agent"
    DESC = "Threat Hunting Agent"

    def __init__(self):
        super().__init__()
        self.analyst_graph: CompiledStateGraph
        self.max_iterations = MAX_ITERATIONS
        self.build_analyst_graph()
        self.build_main_graph()

    def build_analyst_graph(self):
        def analyst_node(state: AnalystState):
            self.logger.debug(f"Analyst Node Invoked (Loop: {state.loop_count})")

            messages = state.messages

            if not messages:
                system_prompt_template = self.load_system_prompt_template("Analyst_System", lang=PROMPT_LANG)
                system_message = system_prompt_template.format()
                human_message = self.load_human_prompt_template("Analyst_Human", lang=PROMPT_LANG).format(
                    question=state.question,
                    case=state.case
                )
                messages = [system_message, human_message]

            llm_api = LLMAPI()

            if state.loop_count >= MAX_ITERATIONS_OF_FUNCTIONS_CALL - 1:
                self.logger.warning("Approaching max iterations, forcing analyst to summarize.")

                stop_instruction = (
                    "\n\n[SYSTEM NOTICE]: You have reached the search limit. "
                    "Do not call any more tools. Please provide your final conclusion "
                    "based ONLY on the information gathered above."
                )

                messages.append(HumanMessage(content=stop_instruction))

                base_llm = llm_api.get_model(tag=["powerful"])
                response: AIMessage = base_llm.invoke(messages)
            else:

                base_llm = llm_api.get_model(tag=["powerful", "function_calling"])
                llm_with_tools = base_llm.bind_tools(tools)
                response: AIMessage = llm_with_tools.invoke(messages)

            if state.loop_count >= MAX_ITERATIONS_OF_FUNCTIONS_CALL - 1:
                if response.tool_calls:
                    self.logger.info("Stripping hallucinated tool calls in final round.")
                    response.tool_calls = []

            return {"loop_count": state.loop_count + 1, "messages": [response]}

        # Tool node
        tool_node = ToolNode(tools)

        def final_answer_node(state: AnalystState):
            self.logger.debug("Final Answer Node Invoked")

            # handle tool_calls
            tool_calls = []
            for message in state.messages:
                if isinstance(message, AIMessage):
                    if message.tool_calls:
                        for tool_call in message.tool_calls:
                            tool_calls.append(tool_call)
                elif isinstance(message, ToolMessage):
                    try:
                        text = json.loads(message.text)
                    except Exception:
                        text = message.text
                    tool_calls.append({"tool_call_id": message.tool_call_id, "name": message.name, "status": message.status, "text": text})
                else:
                    continue

            # get answer reasoning
            last_message = state.messages[-1]

            system_prompt_template = self.load_system_prompt_template("Analyst_Final_System", lang=PROMPT_LANG)
            system_message = system_prompt_template.format()
            human_message = self.load_human_prompt_template("Analyst_Final_Human", lang=PROMPT_LANG).format(
                question=state.question,
                content_to_format=last_message.content
            )
            few_shot_examples = [
            ]
            messages = [
                system_message,
                *few_shot_examples,
                human_message
            ]

            llm_api = LLMAPI()
            formatter_llm = llm_api.get_model(tag=["cheap", "structured_output"])
            structured_llm = formatter_llm.with_structured_output(AnalystOutput)
            response: AnalystOutput = structured_llm.invoke(messages)

            return {
                "answer": response.answer,
                "reasoning": response.reasoning,
                "tool_calls": tool_calls
            }

        # Conditional judgment
        def should_continue(state: AnalystState):
            last_message = state.messages[-1]
            if last_message.tool_calls:
                self.logger.debug("Routing to Tool Node")
                return 'tool'
            self.logger.debug("Routing to Finalizer Node")
            return 'finalizer'

        # --- Build graph ---
        builder = StateGraph(AnalystState)

        builder.add_node('analyst_node', analyst_node)
        builder.add_node('tool', tool_node)
        builder.add_node('finalizer', final_answer_node)

        builder.add_edge(START, 'analyst_node')

        builder.add_conditional_edges(
            'analyst_node',
            should_continue,
        )

        builder.add_edge('tool', 'analyst_node')
        builder.add_edge('finalizer', END)

        self.analyst_graph = builder.compile(name='analyst_graph')

    def build_main_graph(self):
        def intent_node(state: MainState):
            """Intent recognition: determine the overall goal"""
            self.logger.debug("Intent Node Invoked")

            case: CaseModel = Case.get(row_id=self.param_source_row_id)

            user_intent = self.param_user_input

            if not user_intent:
                user_intent = "None (Auto-Pilot Mode)"

            system_prompt_template = self.load_system_prompt_template("Intent_System", lang=PROMPT_LANG)
            system_message = system_prompt_template.format()

            human_message = self.load_human_prompt_template("Intent_Human", lang=PROMPT_LANG).format(
                case=case.model_dump_json_for_ai(profile=AI_PROFILE_INVESTIGATION),
                user_intent=user_intent)

            few_shot_examples = [
            ]

            messages = [
                system_message,
                *few_shot_examples,
                human_message
            ]

            llm_api = LLMAPI()
            llm = llm_api.get_model(tag="fast")
            response: AIMessage = llm.invoke(messages)

            node_out = {
                "case": case,
                "user_intent": user_intent,
                "hunting_objective": response.content,
                "iteration_count": 0,
                "findings": []
            }
            return node_out

        def planner_node(state: MainState):
            """
            Check existing findings to decide what else to look for.
            Generate a batch of tasks at once.
            """
            self.logger.debug("Planner Node Invoked")

            findings = state.findings
            iteration_count = state.iteration_count
            hunting_objective = state.hunting_objective

            iteration_count = iteration_count + 1

            if iteration_count > MAX_ITERATIONS:
                self.logger.debug("Max iterations reached, terminating planning.")
                node_out = {"current_plan": []}
                return node_out

            system_prompt_template = self.load_system_prompt_template("Planner_System", lang=PROMPT_LANG)

            system_message = system_prompt_template.format()

            history_md_list = []
            for record in findings:
                record: Finding
                history_md_list.append(record.to_markdown())

            findings_str = "\n".join(history_md_list)
            additional_info = f"Time Now: {get_current_time_str()}"
            human_message = self.load_human_prompt_template("Planner_Human", lang=PROMPT_LANG).format(case=state.case, hunting_objective=hunting_objective,
                                                                                                      findings=findings_str, iteration_count=iteration_count,
                                                                                                      additional_info=additional_info)

            few_shot_examples = [
            ]

            # Run
            llm_api = LLMAPI()

            llm = llm_api.get_model(tag=["powerful", "structured_output"])

            messages = [
                system_message,
                *few_shot_examples,
                human_message
            ]
            llm = llm.with_structured_output(HuntingPlan)

            response: HuntingPlan = llm.invoke(messages)

            current_record = PlanningRecord(
                iteration=iteration_count,
                rationale=response.rationale,
                plan=response.current_plan
            )
            current_plan = response.current_plan

            self.logger.debug(f"Generated Plan for Round {iteration_count}")

            node_out = {
                "current_plan": current_plan,
                "iteration_count": iteration_count,
                "planning_history": [current_record]

            }
            return node_out

        def continue_to_analysts(state: MainState):
            """
            Conditional edge logic:
            1. If the planner returns a list of tasks -> use the Send API to distribute them to the Subgraph in parallel
            2. If the planner returns an empty list -> end and go to write the report
            """
            current_plan = state.current_plan
            case = state.case
            iteration_count = state.iteration_count
            if not current_plan:
                # No more tasks, end
                self.logger.debug(f"Round {iteration_count},No more tasks in plan, proceeding to report.")
                return "report"

            self.logger.debug(f"Round {iteration_count},Dispatching {len(current_plan)} tasks to analyst subgraph.")
            return [
                Send("analyst_subgraph", AnalystState(question=question, case=case))
                for question in current_plan
            ]

        # --- Encapsulate Subgraph call ---
        def run_analyst_subgraph(state: AnalystState):
            self.logger.debug("Running Analyst Subgraph Wrapper")
            # The output of the graph is dict
            result: dict = self.analyst_graph.invoke(state)
            analyst_state = AnalystState(**result)
            finding = Finding(
                question=analyst_state.question,
                answer=analyst_state.answer,
                reasoning=analyst_state.reasoning,
                tool_calls=analyst_state.tool_calls)

            node_out = {"findings": [finding]}
            return node_out

        def reporter_node(state: MainState):
            """Generate final report"""
            self.logger.debug("Reporter Node Invoked")
            findings = state.findings
            hunting_objective = state.hunting_objective

            # planning_history
            history_md_list = []
            for record in state.planning_history:
                record: PlanningRecord
                history_md_list.append(record.to_markdown())

            planning_history_str = "\n".join(history_md_list)

            # findings
            history_md_list = []
            for record in findings:
                record: Finding
                history_md_list.append(record.to_markdown())

            findings_str = "\n".join(history_md_list)

            # Load system prompt
            system_prompt_template = self.load_system_prompt_template("Report_System", lang=PROMPT_LANG)

            system_message = system_prompt_template.format()
            additional_info = f"Report Time: {get_current_time_str()} \n Reporter: ASF CSIRT Team"
            human_message = self.load_human_prompt_template("Report_Human", lang=PROMPT_LANG).format(hunting_objective=hunting_objective,
                                                                                                     findings=findings_str,
                                                                                                     planning_history=planning_history_str,
                                                                                                     additional_info=additional_info)

            few_shot_examples = [
            ]

            messages = [
                system_message,
                *few_shot_examples,
                human_message
            ]

            llm_api = LLMAPI()
            llm = llm_api.get_model(tag=["powerful"])
            response = llm.invoke(messages)

            case_new = CaseModel(row_id=self.param_source_row_id, threat_hunting_report_ai=response.content)
            Case.update(case_new)

            node_out = {"report": response.content}

            self.update_playbook_status(PlaybookJobStatus.SUCCESS, "Threat Hunting Agent Finish.")
            return node_out

        # --- Build the main graph ---
        main_builder = StateGraph(MainState)

        main_builder.add_node("intent", intent_node)
        main_builder.add_node("planner", planner_node)

        main_builder.add_node("analyst_subgraph", run_analyst_subgraph)
        main_builder.add_node("report", reporter_node)

        main_builder.add_edge(START, "intent")
        main_builder.add_edge("intent", "planner")

        main_builder.add_conditional_edges(
            "planner",
            continue_to_analysts,
            ["analyst_subgraph", "report"]
        )

        main_builder.add_edge("analyst_subgraph", "planner")

        main_builder.add_edge("report", END)

        self.graph = main_builder.compile(checkpointer=self.get_checkpointer())

    def run(self):
        self.run_graph()
        return


if __name__ == "__main__":
    import os
    import django

    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "ASP.settings")
    django.setup()
    model = PlaybookModel(
        source_row_id='141a4bd0-f3cf-4e0c-91b6-f8d9fff6f653',
        user_input="Has the host in the case been infected",
        row_id="401ca83c-4579-4e6f-8329-2e61a6c3405a")
    module = Playbook()
    module._playbook_model = model

    module.run()