# Agent Foundations (TryHackMe)

## Introduction

If you have explored AI applications, AI security, or modern automation workflows, you have probably encountered the word agent. It appears everywhere: research agents, coding agents, SOC agents, browser agents, pentesting agents, customer support agents, and autonomous AI assistants.

However, the term is often used loosely. An “agent” may refer to a simple LLM prompt wrapped in an application, a deterministic workflow controlled by code, or a system that can reason, use tools, maintain state, and decide what to do next. Without clear foundations, it can be difficult to understand what an agent actually is, when one is needed, and which approach is appropriate for building it.

This room is designed to remove that confusion.

Before building the complete **Atlas Research Agent** in the next room, you will work through small, controlled Python examples that introduce the core building blocks of agentic systems. You will begin with a basic LLM workflow, then explore LangChain tool calling, structured outputs, LangGraph state and branching, framework selection, and common debugging failures.

The goal is not to cover every feature of LangChain or LangGraph. Instead, this room focuses on practical engineering judgement: recognising when plain Python is sufficient, when LangChain provides useful abstractions, and when LangGraph is better suited to stateful or branching workflows.

![From LLM Workflow to AI Agent](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786464657156.png align="center")

By the end of the room, you should be able to examine an AI-powered workflow and answer a simple but important question:

**Does this actually need to be an agent?**

## **Learning Objectives**

*   Explain what an AI agent is and when one is needed
    
*   Compare plain Python, LangChain, and LangGraph workflows
    
*   Build a basic Python LLM workflow
    
*   Create a simple LangChain tool-calling workflow
    
*   Use structured outputs to make responses easier to validate
    
*   Build a basic LangGraph workflow with state and branching
    
*   Choose the simplest workflow that safely solves the task
    
*   Debug common agent failures such as invalid tool inputs, missing state updates, and wrong routing decisions
    

## **Prerequisites**

Before starting this room, you should be comfortable with basic Python concepts such as variables, functions, conditionals, loops, files, and imports. If these topics are unfamiliar, complete the [Python Basics](https://tryhackme.com/room/pythonbasics) room first.

Basic familiarity with LLM concepts, including prompts, responses, context, model outputs, and prompt structure, is also recommended. If you need a refresher, the [Prompt Engineering](https://tryhackme.com/room/promptengineeringaisec) room introduces LLM fundamentals, prompt behaviour, and effective prompt design.

## Understanding AI Agents and Framework Choices

![What is an Agent](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786465450691.png align="center")

A common mistake in AI engineering is calling every LLM-powered system an **agent**.

In practice, not every system that uses an LLM needs agentic behaviour. Some workflows only need a prompt, a model, and an output. Others need deterministic code that follows fixed rules. A smaller set of workflows needs the system to make decisions during execution, use tools, track state, and choose what should happen next.

This room separates three ideas that are often confused:

| **Workflow type** | **Description** |
| --- | --- |
| Basic LLM workflow | A fixed path from input to prompt, model, and output |
| Automation | A predictable workflow controlled by code, rules, or functions |
| Agentic workflow | A workflow where the system can choose tools, update state, and decide the next step |

The key difference is **runtime decision-making**. A basic LLM workflow follows a **fixed path**, and although an automated workflow may call functions, its sequence is usually predefined. An agentic workflow introduces greater flexibility by deciding whether a request requires a tool, which tool to use, whether additional information is needed, or which branch should run next. This flexibility can be useful, but it also increases complexity: the more freedom a system has, the harder it becomes to test, debug, evaluate, and secure.

Agentic workflows become more useful when the system must reason over information that is not already available in the prompt. In security, this often means connecting to approved data sources such as endpoint telemetry, identity logs, network events, vulnerability records, previous investigation notes, or ticket history.

Without these connections, the model may produce a plausible answer from incomplete context. Controlled tools allow the workflow to retrieve relevant evidence before deciding what should happen next.

This leads to a simple engineering principle for the room:

**Use the simplest workflow that safely solves the problem.**

## **What Is LangChain?**

[**LangChain**(opens in new tab)](https://www.langchain.com/) is a framework for building LLM-powered applications and agents. In this room, its most important concept is tool calling. A tool is a controlled function that gives the model access to an external capability, such as searching mock data, checking a trusted source, extracting keywords, or classifying a request.

A LangChain agent can decide when to request an approved tool, provide the required input, receive the result, and use that information to produce a final answer.

LangChain is a good fit when a workflow requires:

*   A model call
    
*   A small set of controlled tools
    
*   Simple tool selection
    
*   Clear tool inputs and outputs
    
*   A lightweight agent loop
    

For the early exercises in this room, LangChain introduces tool use without the additional complexity of a full graph-based workflow.

## **What Is LangGraph?**

[**LangGraph**(opens in new tab)](https://www.langchain.com/langgraph) is a framework for building stateful, graph-based agent workflows. Rather than treating the agent as a single loop, it represents the workflow as nodes and edges. Each node performs a specific step, state carries information between steps, edges define the execution path, and conditional edges allow the workflow to branch based on values stored in state.

LangGraph is a good fit when a workflow requires:

*   Explicit state
    
*   Multiple steps
    
*   Conditional routing
    
*   Retries or fallback paths
    
*   Human approval points
    
*   Longer or more complex orchestration
    

For the later exercises in this room, LangGraph makes the execution path visible, testable, and easier to debug.

## **Choosing the Right Approach**

The right implementation depends on the workflow, not on whether the word “agent” sounds more advanced.

*   Use **plain Python** when the task is fixed, predictable, and mostly deterministic.
    
*   Use **LangChain** when the workflow needs an LLM plus a small number of controlled tools, but does not need complex state management or branching.
    
*   Use **LangGraph** when the workflow needs explicit state, conditional paths, retries, human approval points, or structured multi-step orchestration.
    

| **Use case** | **Good fit** |
| --- | --- |
| Fixed prompt and response flow | Plain Python |
| LLM with a few controlled tools | LangChain |
| Stateful workflow with routing or branching | LangGraph |
| Long-running research workflow with synthesis and source handling | Atlas Research Agent |

The purpose of this comparison is not to rank frameworks from simple to advanced. Each option fits a different engineering need.

A simple Python script is often the safest and easiest option when the workflow is predictable. LangChain is useful when a model needs access to controlled tools. LangGraph is better suited for workflows that need visible state, explicit routing, conditional branches, retries, or approval points.

*An agent framework should be selected because the workflow needs it, not because the system uses an LLM.*

![LLM Workflow and Framework Guide](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786465450697.png align="center")

### Answer the questions below

What framework is useful when an LLM needs a few controlled tools? `LangChain`

What type of path does a basic LLM workflow follow? `Fixed path`

## Basic Python LLM Workflow

Before building an agent, it is important to understand the simplest form of an LLM-powered workflow. A basic LLM workflow does not use tools, memory, branching, retries, or stateful orchestration. Instead, it follows a fixed path:

![A left-to-right diagram of a basic LLM workflow. A Python input is turned into a prompt, the prompt is sent to the model, and the model returns an output. The workflow is labelled as a fixed path and shows that tools, memory, state, and branching are not used.](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786465564790.png align="center")

This approach is suitable when the task is predictable, and the system does not need to decide what to do next. In this task, you will run a basic Python script that loads a sample request, constructs a prompt, sends it to the model, and prints the response.

The goal is to understand the foundation on which more advanced workflows are built. LangChain and LangGraph introduce useful abstractions, but the underlying pattern remains the same: provide context to a model and handle its response.

## **Why Start with Plain Python?**

**Plain Python** makes a workflow easier to inspect because it has fewer moving parts: no tools are called, no graph routes execution, and no state is passed between nodes. This makes it easier to understand:

*   What input the workflow receives
    
*   How the prompt is constructed
    
*   What the model returns
    
*   Where errors may occur
    

If a task requires only a single prompt and response, introducing a full agent framework may add unnecessary complexity.

On your machine, open a terminal and move into the lab directory:

Terminal

```powershell
user@machine$ cd ~/agent-foundations
```

## **Working with TODOs**

The lab files contain small `TODO` markers that indicate where you need to complete or modify the code. The surrounding code is already provided so you can focus on one concept at a time.

When you encounter a `TODO`, read the nearby comments first to understand what the missing code should do. After completing it, run the script from the project root and review the output to confirm that your changes work as expected.

In this room, you should only edit the files inside:

```text
agent/
tools/
data/
```

## **The Script**

Open the following file:

Terminal

```powershell
user@machine$ cat agent/01_basic_llm_workflow.py
```

You’ll be able to see that the script performs four simple steps:

1.  Load a sample request.
    
2.  Build a prompt.
    
3.  Send the prompt to the model.
    
4.  Print and log the response.
    

This file has no `TODO` you need to address - it already runs end to end. The only corrections needed are to the example request/prompt and the claimed log output.

The sample request is loaded from `data/sample_requests.json`

Example request:

```json
{
  "request_id": "REQ-001",
  "task": "Summarise the latest research on AI agent tool calling in two sentences.",
  "audience": "junior security analyst",
  "expected_route": "direct_answer"
}
```

The script uses this request to create a prompt similar to:

`You are helping a security team understand an AI research assistant.`  
`Task: Summarise the purpose of an AI research assistant for security analysts.`  
`Audience: junior security analyst`  
`Write a clear and concise explanation.`

The model should return a short response written for the specified **audience**.

## **Run the Workflow**

From the project root, run:

Terminal

```powershell
user@machine$ python3 agent/01_basic_llm_workflow.py
```

Review the output in the terminal, then check the log file `logs/agent_foundations.log`.

Each line is formatted as `timestamp | LEVEL | logger name | message`. You should see these messages, in order:

`Loaded request REQ-001`  
`'Summarise the latest research on AI agent tool calling in two sentences.'`  
`Built prompt from request.`  
`Model response received.`  
`Workflow run completed.`

## **Modify the Request**

Open the `data/sample_requests.json` file:

Terminal

```powershell
user@machine$ nano data/sample_requests.json
```

Find the request with `request_id` `REQ-001` and change the audience from:

`junior security analyst`

To:

`SOC manager`

Run the script again:

Terminal

```powershell
user@machine$ python3 agent/01_basic_llm_workflow.py
```

Then compare the new response with the previous one. The main topic should remain consistent, but the explanation should adapt to the new audience.

After completing this task, you should be able to explain how a basic Python LLM workflow operates and why it is not yet an agent. The workflow does not select tools, update state, branch between paths, retry failed steps, or decide what action to take next. Instead, it follows a fixed execution path from input to output.

### Answer the questions below

Does this workflow use tools? (Yea/Nay) `Nay`

What changes the response style? `Audience`

Is this an agent? (Yea/Nay) `Nay`

Run the basic workflow after modifying the audience for REQ-001. What is the flag?

```shell
python3 agent/01_basic_llm_workflow.py
2026-08-25 09:18:14 | INFO     | config | Using the THM.
=== 01_basic_llm_workflow: Input -> Prompt -> Model -> Output ===
2026-08-25 09:18:14 | INFO     | 01_basic_llm_workflow | Loaded request REQ-001 for audience 'SOC manager' - 'Summarise the latest research on AI agent tool calling in two sentences.'
2026-08-25 09:18:14 | INFO     | 01_basic_llm_workflow | Built prompt from request.
2026-08-25 09:18:17 | INFO     | 01_basic_llm_workflow | Model response received.

Model response:
The latest research on AI agent tool calling highlights the ability of AI systems to autonomously invoke external tools and APIs, enhancing their operational capabilities in real-time decision-making. This advancement allows for more efficient incident response and threat detection by enabling agents to access and utilize relevant data and functionalities as needed.
2026-08-25 09:18:17 | INFO     | 01_basic_llm_workflow | Workflow run complete.
2026-08-25 09:18:17 | INFO     | 01_basic_llm_workflow | THM{basic_llm_workflow_REDACTED}

THM{basic_llm_workflow_REDACTED}
```

## Tool Calling with LangChain

In the previous task, you ran a basic LLM workflow that followed a fixed path:

Input → Prompt → Model → Output

This approach is useful for predictable tasks, but it has an important limitation: the model can respond only using the context provided in the prompt and the knowledge available to it.

The image below compares a fixed LLM workflow with a small LangChain-style tool-calling agent. The key difference is that tool calling allows the model to request approved external capabilities instead of relying only on prompt context and model knowledge.

![Basic workflow and Langchain](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786465966103.png align="center")

In this task, you will build a small LangChain tool-calling agent. Rather than creating the full **Atlas Research Agent** yet, you will focus on the core interaction between the user, the model, an approved tool, the tool result, and the final response.

## **Why Connect Tools to Data Sources?**

A model does not automatically know what is happening inside an environment. In security workflows, relevant evidence may be distributed across endpoint telemetry, identity logs, network events, cloud alerts, vulnerability records, previous investigation notes, and ticket history.

Each platform may provide strong visibility into one part of the environment without capturing the complete picture. For example, an endpoint tool may reveal process and device activity, while identity, network, cloud, or historical investigation context remains in other systems.

Tool calling helps bridge these gaps by connecting the workflow to approved data sources rather than forcing the model to rely on incomplete context or guess. In a well-designed agent, tools are therefore not merely additional capabilities; they are controlled windows into trusted data.

![Tools and richer investigation view](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786974702240.png align="center")

## **Why Use LangChain?**

Tool calling can be implemented manually, but frameworks such as **LangChain** provide standard interfaces for connecting language models, prompts, and tools. LangChain is useful when an application needs a small set of controlled tools and a simple execution loop, allowing the model to request a tool, provide its arguments, receive the result, and continue the workflow.

However, LangChain does not make a system safe by default. The developer must still control which tools are available, validate inputs, handle errors, and restrict what each tool is allowed to do. The model may decide which approved tool to request, but the Python application remains responsible for validating and executing that request.

## **Defining Tools**

On your machine, open the following file:

```shell
user@machine$ nano ~/agent-foundations/tools/source_tools.py
```

The file already has its imports, `ALLOWED_SEVERITIES`, the `CVE_ID_PATTERN` regex, and a mock `SECURITY_RECORDS` dict keyed by CVE ID. It also already has both tools stubbed out with their `@tool` decorator, docstring, and a `TODO` where the lookup logic goes:

```python
@tool
def search_security_record(cve_id: str) -> dict:
    """
    Look up a single security record by its CVE identifier (e.g. 'CVE-2025-1001').

    Returns a dict with a "found" key. When found, also includes "cve_id",
    "title", "severity", and "description". When not found - including when
    cve_id isn't a validly formatted CVE identifier - includes an "error"
    key instead.
    """
    # TODO: validate `cve_id` against CVE_ID_PATTERN, then look it up in
    # SECURITY_RECORDS.
```

Fill in the `TODO`:

```plaintext
@tool
def search_security_record(cve_id: str) -> dict:
    """
    Look up a single security record by its CVE identifier (e.g. 'CVE-2025-1001').

    Returns a dict with a "found" key. When found, also includes "cve_id",
    "title", "severity", and "description". When not found - including when
    cve_id isn't a validly formatted CVE identifier - includes an "error"
    key instead.
    """
    if not isinstance(cve_id, str) or not CVE_ID_PATTERN.match(cve_id.strip()):
        return {
            "found": False,
            "cve_id": str(cve_id),
            "error": "Invalid CVE identifier format. Expected 'CVE-YYYY-NNNN'.",
        }

    normalized = cve_id.strip().upper()
    record = SECURITY_RECORDS.get(normalized)
    if record is None:
        return {
            "found": False,
            "cve_id": normalized,
            "error": "No security record found for this CVE identifier.",
        }

    return {"found": True, "cve_id": normalized, **record}
```

Do the same for the second tool:

```plaintext
@tool
def list_security_records(severity: Literal["critical", "high", "medium", "low"]) -> dict:
    """
    List security records filtered by severity.

    Args:
        severity: One of 'critical', 'high', 'medium', 'low'.

    Returns:
        {"severity_filter": str, "count": int,
         "records": [{"cve_id": str, "title": str, "severity": str}, ...]}
    """
    if severity not in ALLOWED_SEVERITIES:
        return {
            "severity_filter": str(severity),
            "count": 0,
            "records": [],
            "error": f"severity must be one of {ALLOWED_SEVERITIES}.",
        }

    records = [
        {"cve_id": cve_id, "title": data["title"], "severity": data["severity"]}
        for cve_id, data in SECURITY_RECORDS.items()
        if data["severity"] == severity
    ]

    return {"severity_filter": severity, "count": len(records), "records": records}
```

The `@tool` decorator converts each Python function into a LangChain tool. The function name, docstring, and type annotations describe the tool to the model, helping it determine when the tool should be used and which arguments it requires. Each tool returns a structured dictionary rather than free-form text, making the result easier to inspect, validate, and reuse throughout the workflow. Both tools validate their input against `SECURITY_RECORDS/ALLOWED_SEVERITIES` before doing anything else - never trust that the model's argument is well-formed.

### Building the Agent Open the agent file:

```python
user@machine$ nano agent/02_langchain_tool_agent.py
```

Add the required import:

```python
from langchain.agents import create_agent
```

The model and the approved tool set are already defined:

```python
model = get_model(supports_tools=True)
tools = [search_security_record, list_security_records]
```

The model and the approved tool set are already defined:

```python
model = get_model(supports_tools=True)
tools = [search_security_record, list_security_records]
```

Now replace the `agent = None` placeholder with a real agent:

```python
agent = create_agent(
    model=model,
    tools=tools,
    system_prompt=(
        "You are a security research assistant. "
        "Use only the provided tools to retrieve security records. "
        "Do not invent records or tool results. "
        "Clearly state when a record cannot be found."
    ),
)
```

The system prompt defines the agent's role and boundaries. It instructs the model to rely on approved tools instead of inventing security records. `run_agent()` is already implemented - it invokes the agent with the user's request and returns the content of the final message:

```plaintext
def run_agent(user_request: str) -> str:
    """
    Invoke the agent with a single user message and return the final
    assistant message's content.
    """
    if agent is None:
        # Safe placeholder: lets this script run before the agent above is built.
        logger.info("run_agent() is using the placeholder implementation (agent is None).")
        return "TODO: run_agent() is not implemented yet."

    result = agent.invoke({"messages": [{"role": "user", "content": user_request}]})
    return result["messages"][-1].content
```

Once `agent` is a real `create_agent(...)` object instead of `None`, the placeholder branch above never runs.

The entry point is also already written, and prints the flag once the agent's response actually mentions the CVE it was asked about:

```plaintext
if __name__ == "__main__":
    print("=== 02_langchain_tool_agent ===")

    response = run_agent("Find the security record for CVE-2025-1001.")
    print(response)

    if agent is not None and "TODO" not in response and "CVE-2025-1001" in response:
        flag = get_flag("langchain_tool_agent_working")
        logger.info(flag)
        print(f"\n{flag}")
    else:
        print("\nThe agent isn't fully working yet - no flag.")
```

### Run the Agent

```python
user@machine$ python3 agent/02_langchain_tool_agent.py
```

The agent should use the approved tool to retrieve the record for CVE-2025-1001, produce a short response based on the tool result, and print the flag.

Now temporarily change the request passed to `run_agent()` in the entry point in `02_langchain_tool_agent.py` to:

```python
response = run_agent("Find the security record for 1001.")
```

Run the script again:

```python
user@machine$ python3 agent/02_langchain_tool_agent.py
```

`search_security_record` should reject `"1001"` as an invalid CVE format and the agent should report that no record could be found, rather than fabricating one - `"CVE-2025-1001"` won't appear in the response, so the flag won't print for this run. That's expected: it demonstrates that the tool doesn't blindly trust the model's argument, but validates the input before processing it and returning a result.

Change the request back to:

```python
response = run_agent("Find the security record for CVE-2025-1001.")
```

Before moving on, so the script prints the flag again.

### Answer the questions below

What decorator converts a Python function into a LangChain tool? `@tool`

Modify then run the LangChain tool script. What is the flag THM{langchain\_tool\_agent\_REDACTED}

```python
nano agent/02_langchain_tool_agent.py

ubuntu@tryhackme:~/agent-foundations$ python3 agent/02_langchain_tool_agent.py
2026-08-26 13:12:51 | INFO     | config | Using the THM.
=== 02_langchain_tool_agent ===
CVE-2025-1001 (critical severity): Remote code execution in mock-http-server via crafted header. A crafted request header allows an unauthenticated attacker to trigger remote code execution on affected mock-http-server deployments prior to version 2.3.1.
2026-08-26 13:12:51 | INFO     | 02_langchain_tool_agent | THM{langchain_tool_agent_REDACTED}

THM{langchain_tool_agent_REDACTED}
```

## Structured Outputs

In the previous task, the LangChain agent returned a natural-language response. Although this format is useful for humans, it can be difficult for software to validate, test, or reuse. Agent workflows often require predictable, structured outputs so that other parts of the application can reliably extract information such as:

*   Was a tool needed?
    
*   What risk level was assigned?
    
*   What should happen next?
    
*   Was the answer valid?
    

If the model returns a free-form paragraph, important details may be missing, inconsistent, or difficult for software to parse. Structured outputs address this problem by requiring the response to follow a defined schema.

Instead of returning only text, the workflow can produce fields such as:

*   `request_id`
    
*   `summary`
    
*   `risk_level`
    

This makes the output easier to validate, log, test, reuse, and pass to the next step in the workflow.

## **Why Structured Outputs Matter**

Free-form model responses are flexible, but that flexibility can create ambiguity. For example, a model might return:

*This request looks important. I think the analyst should review it soon.*

Although the response sounds reasonable, it does not clearly indicate the assigned risk level, whether a tool was required, or what action should happen next.

A structured output is clearer:

```json
{ 
 "request_id": "REQ-003",
 "summary": "The request asks for a security record review.",
 "needs_tools": true,
 "risk_level": "medium",
 "next_action": "retrieve_security_record",
 "confidence": 0.78 
}
```

The second response is easier to check because each field has a specific purpose.

![Structured Outputs](https://cdn-images.tryhackme.com/user-uploads/5f5ed9259575d24307292950/room-content/5f5ed9259575d24307292950-1786972844367.png align="center")

## **The Script**

On your machine, open the following file:

```python
user@machine$ nano agent/03_structured_outputs.py
```

This script already defines the output contract for you: a `StructuredResponse` Pydantic model with `request_id`, `summary`, `needs_tools`, `risk_level`, `next_action`, and `confidence` fields, plus validators that reject anything outside the allowed values:

```python
ALLOWED_RISK_LEVELS = ["low", "medium", "high"]
ALLOWED_NEXT_ACTIONS = [
    "answer_directly",
    "retrieve_sources",
    "request_human_review",
    "reject_request",
]

class StructuredResponse(BaseModel):
    """The schema every structured agent response must satisfy."""

    request_id: str
    summary: str
    needs_tools: bool
    risk_level: str
    next_action: str
    confidence: float

    @field_validator("risk_level")
    @classmethod
    def validate_risk_level(cls, value: str) -> str:
        if value not in ALLOWED_RISK_LEVELS:
            raise ValueError(f"risk_level must be one of {ALLOWED_RISK_LEVELS}, got {value!r}")
        return value

    @field_validator("next_action")
    @classmethod
    def validate_next_action(cls, value: str) -> str:
        if value not in ALLOWED_NEXT_ACTIONS:
            raise ValueError(f"next_action must be one of {ALLOWED_NEXT_ACTIONS}, got {value!r}")
        return value

    @field_validator("confidence")
    @classmethod
    def validate_confidence(cls, value: float) -> float:
        if not (0.0 <= value <= 1.0):
            raise ValueError(f"confidence must be between 0.0 and 1.0, got {value!r}")
        return value
```

Restricting `risk_level` and `next_action` to a fixed set of values makes invalid or inconsistent model output easy to catch before anything downstream trusts it.

### Filling In the Classification Logic

The schema is already in place - your job is the `TODO` inside `generate_structured_response()`.

This function currently builds a placeholder candidate that always reports `low` risk and `answer_directly`, regardless of the actual request. Replace the placeholder with logic that calls the shared `classify_request_type()` helper (`from agent/config.py`) and maps its result to `risk_level`, `next_action`, and `needs_tools`:

![](https://cdn.hashnode.com/uploads/covers/5f4a98085ee1ba597542e097/57b30c83-c413-47c5-b25b-276e51328137.png align="center")

`confidence` should follow `{"low": 0.9, "medium": 0.75, "high": 0.55}[risk_level]`.

```plaintext
def generate_structured_response(request: dict) -> dict:
    """Build a candidate structured response from a classification + model summary."""
    task_text = request.get("task", "")

    model = get_model()
    model_summary = model.invoke(f"Summarise this request for a structured report: {task_text}").content
    logger.info("Model produced a raw summary for the structured response.")

    classification = classify_request_type(task_text)
    request_type = classification["request_type"]

    if request_type == "invalid":
        risk_level, next_action, needs_tools = "low", "reject_request", False
    elif request_type == "human_review":
        risk_level, next_action, needs_tools = "high", "request_human_review", False
    elif request_type == "source_lookup":
        risk_level, next_action, needs_tools = "medium", "retrieve_sources", True
    else:  # direct_answer
        risk_level, next_action, needs_tools = "low", "answer_directly", False

    confidence = {"low": 0.9, "medium": 0.75, "high": 0.55}[risk_level]

    candidate = {
        "request_id": request["request_id"],
        "summary": model_summary,
        "needs_tools": needs_tools,
        "risk_level": risk_level,
        "next_action": next_action,
        "confidence": confidence,
    }
    logger.info(f"Structured candidate built for {request['request_id']}: {candidate}")
    return candidate
```

The rest of the file (`load_sample_request`, `validate_structured_response`, and the `__main__` block) is already complete - it validates your candidate against `StructuredResponse` and prints the flag once `generate_structured_response()` correctly classifies the sample request (`REQ-002`, a source-lookup request) as `next_action="retrieve_sources"`, `needs_tools=True`, `risk_level="medium"`.

### Run the Script

The entry point at the bottom of the file is already written for you - you don't need to add one:

```plaintext
if __name__ == "__main__":
    print("=== 03_structured_outputs: generate -> validate -> emit structured JSON ===")

    sample_request = load_sample_request("REQ-002")
    candidate_response = generate_structured_response(sample_request)

    validated_response = validate_structured_response(candidate_response)

    print("\nValidated structured output:")
    print(json.dumps(validated_response.model_dump(), indent=2))

    if (
        validated_response.next_action == "retrieve_sources"
        and validated_response.needs_tools is True
        and validated_response.risk_level == "medium"
    ):
        flag = get_flag("structured_outputs_validated")
        logger.info(flag)
        print(f"\n{flag}")
    else:
        print("\ngenerate_structured_response() isn't using classify_request_type() yet - no flag.")
```

Run it:

```python
user@machine$ python3 agent/03_structured_outputs.py
```

Once your mapping is in place, you should see a JSON response that follows the `StructuredResponse` schema for `REQ-002` (a source-lookup request), and the flag printed.

### Test the Validation

Now break it on purpose. In generate\_structured\_response(), temporarily change the source\_lookup branch from:

```python
elif request_type == "source_lookup": risk_level, next_action, needs_tools = "medium", "retrieve_sources", True 
```

To:

```python
elif request_type == "source_lookup": risk_level, next_action, needs_tools = "critical", "retrieve_sources", True
```

Run the script again:

```python
Terminal user@machine$ python3 agent/03_structured_outputs.py 
```

`validate_structured_response()` should raise a `ValidationError`, because `"critical"` is not one of `ALLOWED_RISK_LEVELS`. This demonstrates how the schema's validators catch an invalid value before anything downstream trusts it.

```python
elif request_type == "source_lookup": risk_level, next_action, needs_tools = "medium", "retrieve_sources", True
```

### Answer the questions below

What Pydantic model validates the structured output? `StructuredResponse`

## State with LangGraph

In the previous tasks, the workflow either followed a fixed path or returned a structured decision. Real agent workflows, however, often need to preserve information across multiple steps, such as the original request, its type, whether a tool was used, whether an error occurred, and the final answer.

This tracked information is called state. By carrying a structured object through the workflow instead of passing loose values between functions, state makes each step easier to inspect, test, and debug.

In this task, you will use LangGraph to build a small stateful workflow that follows this path:

START → load\_request → classify\_request → generate\_answer → END

Each node will read from the current state, update the relevant fields, and return only the fields it changed - LangGraph merges that into the shared state before handing it to the next node.

## **Why State Matters**

A basic Python workflow can pass values directly between functions, which works well for small scripts. As workflows become more agentic, however, they often need to track additional information:

*   What was the original request?
    
*   What type of request is it?
    
*   Were tools required?
    
*   Did any step fail?
    
*   What answer was generated?
    

State provides a shared structure for storing and updating information throughout the workflow. Because each step records what it changed, a stateful workflow is easier to inspect and debug, especially before introducing branching, retries, or human approval points.

This becomes even more important when the workflow connects to multiple data sources. If an agent checks endpoint telemetry, identity logs, vulnerability records, and previous notes, it should track which sources were queried, what each source returned, and which evidence is still missing. Without this shared state, it becomes much harder to explain how the agent reached its conclusion.

## **The Script**

Open the following file:

Terminal

```powershell
user@machine$ nano agent/04_langgraph_state.py
```

The state schema is already defined for you:

```python
class AgentState(TypedDict):
    """The shared state object every node in the graph reads and updates."""

    request_id: str
    user_request: str
    request_type: str
    tool_results: list
    errors: list
    final_answer: str
```

`AgentState` defines the fields that move through the graph:

*   `request_id`
    
*   `user_request`
    
*   `request_type`
    
*   `tool_results`
    
*   `errors`
    
*   `final_answer`
    

### The Nodes That Are Already Wired

A LangGraph node is a function that receives the current state and returns a dict of the fields it wants to update - it doesn't need to return the whole state back.

Three of the four nodes are already implemented for you:

```plaintext
def load_request(state: AgentState) -> dict:
    """Node 1: populate request_id and user_request from sample data."""
    with open(SAMPLE_REQUESTS_PATH, "r", encoding="utf-8") as f:
        requests = json.load(f)

    request = requests[0]
    logger.info(f"[load_request] Loaded {request['request_id']}.")

    return {
        "request_id": request["request_id"],
        "user_request": request["task"],
    }


def classify_request(state: AgentState) -> dict:
    """Node 2: classify the request and record the result in tool_results."""
    classification = classify_request_type(state["user_request"])
    logger.info(f"[classify_request] request_type={classification['request_type']}")

    updated_tool_results = state["tool_results"] + [
        {"tool": "classify_request_type", "result": classification}
    ]

    return {
        "request_type": classification["request_type"],
        "tool_results": updated_tool_results,
    }


def log_summary(state: AgentState) -> dict:
    """Node 3: log a summary of state so far. No state changes needed."""
    logger.info(
        f"[log_summary] request_id={state['request_id']} "
        f"request_type={state['request_type']} "
        f"tool_calls_so_far={len(state['tool_results'])}"
    )
    return {}
```

`load_request` reads the first sample request from `data/sample_requests.json`. `classify_request` uses `classify_request_type()` from `agent/config.py`, the same helper introduced in the previous task. It also records the call in `tool_results` so the state shows which "tools" ran. `log_summary` doesn't change any state - it just logs where the workflow is so far.

### Finishing the Fourth

Node Your first job is the TODO in `generate_answer()`. Right now, the function returns a placeholder string.

Replace it with logic that uses `state["request_type"]` and `state["user_request"]` to build a prompt, then calls the model.

```plaintext
def generate_answer(state: AgentState) -> dict:
    """Node 4: call the model to produce the final answer."""
    prompt = (
        f"Summarise this {state['request_type']} request in one sentence: "
        f"{state['user_request']}"
    )
    answer = get_model().invoke(prompt).content
    logger.info("[generate_answer] Final answer generated.")

    return {"final_answer": answer}
```

### Building the Graph

Your second job is the `TODO` in `build_graph()`. It currently only registers `load_request` and wires it straight to `END` as a safe placeholder. Register the other three nodes and connect all five edges:

```plaintext
def build_graph() -> StateGraph:
    """Wire the nodes together into the START -> ... -> END graph."""
    graph = StateGraph(AgentState)

    graph.add_node("load_request", load_request)
    graph.add_node("classify_request", classify_request)
    graph.add_node("log_summary", log_summary)
    graph.add_node("generate_answer", generate_answer)

    graph.add_edge(START, "load_request")
    graph.add_edge("load_request", "classify_request")
    graph.add_edge("classify_request", "log_summary")
    graph.add_edge("log_summary", "generate_answer")
    graph.add_edge("generate_answer", END)

    return graph.compile()
```

This defines the execution path: the graph starts, loads the request, classifies it, logs a summary, generates an answer, and ends.

### Run the Workflow

The entry point at the bottom of the file is already written for you:

```plaintext
if __name__ == "__main__":
    print("=== 04_langgraph_state: START -> load_request -> classify_request -> log_summary -> generate_answer -> END ===")

    compiled_graph = build_graph()

    initial_state: AgentState = {
        "request_id": "",
        "user_request": "",
        "request_type": "",
        "tool_results": [],
        "errors": [],
        "final_answer": "",
    }

    final_state = compiled_graph.invoke(initial_state)

    print("\nFinal state:")
    print(json.dumps(final_state, indent=2))

    logger.info("LangGraph run complete.")

    if (
        final_state["request_type"]
        and final_state["final_answer"]
        and "TODO" not in final_state["final_answer"]
    ):
        flag = get_flag("langgraph_nodes_connected")
        logger.info(flag)
        print(f"\n{flag}")
    else:
        print("\nThe graph isn't fully wired yet (build_graph()/generate_answer()) - no flag.")
```

Run it:

```python
user@machine$ python3 agent/04_langgraph_state.py
```

The final state should be printed in the terminal, along with the flag once both `TODO` are done. Notice that the workflow returns more than an answer: it also exposes the information carried and updated throughout the graph (`request_type`, `tool_results`, and so on).

### Answer the questions below

What class defines the workflow state? `AgentState`

## Branching, Routing, and Framework Decision

In the previous task, you used LangGraph state to carry information through a fixed workflow:

START → load\_request → classify\_request → generate\_answer → END

This approach is useful, but many agent workflows need to choose different paths depending on the request. A simple question may require a direct answer, a CVE-related request may need a source lookup, and an unsafe or unsupported request may need to be rejected.

This is where branching and routing become useful. Branching allows the workflow to select the next node based on values stored in state, so each request can follow the path that matches its needs.

In this task, you will complete a LangGraph workflow with conditional routing. Every request begins at `classify_request` and then passes through route\_request, which sends it to exactly one of four branches: `direct_answer, retrieve_sources, needs_human_review,` or `reject_invalid_request`. After the selected branch completes, the workflow rejoins at a shared finalize node before reaching the end of the graph.

The request follows one of four branches: direct\_answer, retrieve\_sources, needs\_human\_review, or reject\_invalid\_request. All four rejoin at a common finalize node before the graph ends.

The workflow stores important fields in state - request\_type, route, tool\_results, errors, final\_answer - and route\_request uses request\_type to decide which of the four branches to take. This is why state becomes especially useful once workflows start branching: it gives the graph a visible record of what happened and why a route was selected.

State with langgraph

This is why state becomes especially useful when workflows start branching. It gives the graph a visible record of what happened and why a route was selected.

### Why Branching Matters

A basic workflow is easy to understand because every request follows the same path. Security and research workflows, however, often require different behaviour depending on the request, such as answering directly, retrieving evidence, stopping safely, or rejecting something outside the workflow’s scope.

Branching makes this decision explicit. Rather than hiding routing logic inside a long prompt, the workflow stores the request type in state and uses that value to select the next step. This makes the execution path easier to inspect, test, and debug.

### The Script

Open the following file:

```python
user@machine$ nano agent/05_langgraph_branching.py
```

The state and the routing table are already defined for you:

```python
ALLOWED_ROUTES = [
    "direct_answer",
    "retrieve_sources",
    "needs_human_review",
    "reject_invalid_request",
]

# Maps classify_request_type()'s request_type onto this graph's branch names.
_ROUTE_BY_REQUEST_TYPE = {
    "direct_answer": "direct_answer",
    "source_lookup": "retrieve_sources",
    "human_review": "needs_human_review",
    "invalid": "reject_invalid_request",
}


class BranchingState(TypedDict):
    request_id: str
    user_request: str
    request_type: str
    route: str
    tool_results: list
    errors: list
    final_answer: str
```

`request_type` is the value the request classifier assigns; route is the branch name `route_request` decides on. `_ROUTE_BY_REQUEST_TYPE` is the table that connects the two.

### The Nodes That Are Already Wired

`classify_request` is already implemented - it reuses the same `classify_request_type()` helper from `agent/config.py` you've used in the last two tasks, and records the call in `tool_results`:

```python
def classify_request(state: BranchingState) -> dict:
    """Classify the request so route_request has something to branch on."""
    classification = classify_request_type(state["user_request"])
    logger.info(f"[classify_request] request_type={classification['request_type']}")
    return {
        "request_type": classification["request_type"],
        "tool_results": state["tool_results"] + [
            {"tool": "classify_request_type", "result": classification}
        ],
    }
```

All four branch nodes, and the `finalize` node they all rejoin at, are also already implemented:

```python
def direct_answer(state: BranchingState) -> dict:
    """Branch: answer simple requests directly with the model."""
    prompt = f"Answer this request directly: {state['user_request']}"
    answer = get_model().invoke(prompt).content
    logger.info("[direct_answer] Answered directly, no source retrieval needed.")
    return {"final_answer": answer}


def retrieve_sources(state: BranchingState) -> dict:
    """Branch: retrieve security records before answering."""
    records = list_security_records.invoke({"severity": "critical"})
    logger.info(f"[retrieve_sources] Retrieved {records['count']} critical record(s).")

    answer = f"Found {records['count']} critical security record(s) to cite in this answer."
    return {
        "final_answer": answer,
        "tool_results": state["tool_results"] + [
            {"tool": "list_security_records", "result": records}
        ],
    }


def needs_human_review(state: BranchingState) -> dict:
    """Branch: stop and flag the request for a human instead of answering."""
    logger.info("[needs_human_review] Request flagged for human review; no automatic answer given.")
    return {
        "final_answer": "This request requires human review before the agent can proceed.",
    }


def reject_invalid_request(state: BranchingState) -> dict:
    """Branch: reject requests that are empty or otherwise invalid."""
    logger.warning("[reject_invalid_request] Request rejected as invalid.")
    return {
        "final_answer": "This request could not be processed - it was empty or invalid.",
        "errors": state["errors"] + ["invalid_request"],
    }


def finalize(state: BranchingState) -> dict:
    """Common exit node: log the final outcome for every branch."""
    logger.info(f"[finalize] route={state['route']} final_answer={state['final_answer']!r}")
    return {}
```

Notice `retrieve_sources` calls the real `list_security_records` tool from `tools/source_`[`tools.py`](http://tools.py) - the same tool file you completed earlier in this room - instead of a hardcoded answer. This is also where `needs_human_review` comes in: it's a fourth branch that stops the workflow and asks for a human instead of generating an answer, for requests `classify_request_type()` flags as `human_review`.

The conditional-edge selector is already implemented too:

```python
def select_branch(state: BranchingState) -> str:
    """Conditional-edge function: reads state['route'] and returns the next node name."""
    route = state["route"]
    if route not in ALLOWED_ROUTES:
        logger.warning(f"[select_branch] Unknown route {route!r}; defaulting to reject_invalid_request.")
        return "reject_invalid_request"
    return route
```

## Debugging Common Agent Failures

## Conclusion