Guardrail Timing issue - Code Review Assignment CIM2_Assignment.ipynb

[kashyap3881]

In this Code Review Crew consisting of senior_developer, security_engineer, and tech_lead agents, a guardrail was added to the review_security task associated with the security_engineer Agent as shown below. What was observed was both strange and interesting (please see picture below) where you will see that even though the security_engineer produced a clear and good json output the guardrail blocked it as json_dict is None ! Even though the guardrail was triggered after the Security Engineer final answer, it still identified json_dict as None, which tells me that there’s timing issue with how CrewAI processes the output and when the guardrail runs.

1. The Security Engineer agent produces a valid JSON response (shown in the output figure)
2. The guardrail function security_review_output_guardrail is called
3. In security_review_output_guardrail function json_output = output if type(output)==dict else output.json_dict
4. The output object is a TaskOutput object, not a dict
5. So it tries to access output.json_dict
6. However, it looks like at the moment the guardrail runs, the json_dict attribute might not be populated yet (it’s None)
7. This is because CrewAI might be parsing the LLM’s text output into the Pydantic model after the guardrail runs, or the parsing hasn’t completed yet ?

How could we ensure that the guardrails are triggered smoothly by CrewAI after it has done the required Pydantic mode parsing ? The next trigger of the security engineer to re-evaluate was unnecessary and was caused by this timing issue. Fixing this will not only make the responses faster but also save the cost of re-triggering the Agent to evaluate and use the LLM for security analysis.

Task with guardrail———

class SecurityVulnerability(BaseModel):
    description: str
    risk_level: str
    evidence: str

# Define the pydantic model for the security review output
class ReviewSecurityJSON(BaseModel):
    security_vulnerabilities: list[SecurityVulnerability] 
    blocking: bool
    highest_risk: str
    security_recommendations: list[str]

# Create the security review task
review_security = Task(
    # Load the expected output, agent, and name from the YAML configuration
    config=tasks_config['review_security'], 
    # Define the output type as a pydantic model
    output_json=ReviewSecurityJSON,
    # Add the security_review_output_guardrail guardrail
    guardrails=[security_review_output_guardrail],     -----> guardrail before going to the tech_lead agent
    # Define the agent that will perform this task
    agent=security_engineer 
)

[magdalena.bouza]

Hey @kashyap3881. Guardrails are supposed to be called on the Task Output, so that shouldn’t be the issue. My guess is that the task is returning a dictionary, rather than a JSON. Please make sure that the Task definition expected output mentions a JSON, and not a dictionary.

[kashyap3881]

Hi @magdalena.bouza,

Thank you for the response. Kindly note guardrail “security_review_output_guardrail” is applied in review_security Task as shown in my earlier message (I also added it below) and the output is a pydantic structure with clear definition and the output_json is also set in the task as “output_json=ReviewSecurityJSON”.
The security_review_output_guardrail function also checks for the dictionary (code shown below and we use clear pydantic structure so this is not a problem) and its clear that it is receiving None. The task is always generating a valid JSON (all the time) and the fact the guardrail is detecting None means that the json_dict attribute might not be populated yet when Crew AI is calling the guardrail.
Its clearly a timing issue and is easily reproducible. The second time the guardrail passes but that does not mean the timing issue does not exist since it would be using the json_dict in the memory as i set memory to True (please see the Crew code below).

Versions used:
crewai: 1.7.2
crewai_tools: 1.7.2
Pyhton version: 3.12.2

Review Security task Code:

# Removed code from post

Security Review Output Guardrail Function Code:

# GRADED CELL: Exercise 1

def security_review_output_guardrail(output):
    
    # get the (JSON) output from the TaskOutput object
    try:
        json_output = output if type(output)==dict else output.json_dict -------> checks if output is dict 
    except Exception as e:
        return (False, ("Error retrieving the `json_dict` argument: "
                        f"\n{str(e)}\n"
                        "Make sure you set the output_json parameter in the Task."
                        )
                )
    
    if json_output is None:
        return (False, "The json_dict is None. Make sure the task is producing valid JSON output.")

    # define risk levels
    valid_risk_levels = ['low', 'medium', 'high']

    # ... rest of the validation logic

Crew Memory set to True:

# Removed code from post

[magdalena.bouza]

Hi @kashyap3881 I’m really sorry, I had to edit your post because it contained solutions for the assignment, and that’s not allowed, but I did look carefully at what you showed there.

I wanted to point out one detail in you process: the beginning of the guardrail is only there to make unittests easier to write. During execution, the Guardail will receive a TaskOutput object, so the try/except clause is only there to see if it is a dummy dictionary for testing or an actual output from a Task, it doesn’t catch the scenario where you are producing a result in the form of a dictionary, rather than a JSON. My guess is still that you are getting the first output as a dictionary rather than an actual JSON (the structure in the final answer would look the same). Since you are seeing the error that the json_dict is None, the guardrail is receiving the TaskOutput object, which is generated when the task completes.

[kashyap3881]

Hi @magdalena.bouza,

The code shared in security_review_output_guardrail was not a part of the solution, since solutions exist between ‘### START CODE HERE ###’ and ‘### END CODE HERE ###’ and I did not share my solution for it.

The code shared for the crew was my solution and I appreciate you removing that. It is pretty straightforward to simply listing all the tasks, agents, before_kickoff_callbacks, etc.

I had written my custom callback and removed the guardrail (will explain later why i had to remove the guardrail) and verified that every single output from the Security Engineer Agent is a valid JSON. I have repeated this multiple times and there should be no doubt whatsoever that the Security Engineer Agent is producing anything else other than a valid JSON. The strict pydantic data structures help us to achieve this and the output_json is set.
So this is a clear case of timing problem because when I add back the guardrail the guardrail is simply called by the CrewAI with None json_dict.

Why I had to remove the guardrail for my investigation is because of the flow below. We need to have something that allows us to debug/check output formats of the TASK BEFORE the guardrail is called and I believe CrewAI does not have any feature at the moment that supports this.

Task Agent Produces Output
↓
Guardrail Runs ← json_dict is None here! (timing issue)
↓
[Retry if guardrail fails]
↓
Guardrail Passes
↓
Callback Runs ← Only runs AFTER guardrails complete

I will investigate the source code in GitHub - crewAIInc/crewAI at 1.7.2 to see what is causing this once I am back from my break and will open the issue.
Also may be a new debug feature to allow checks for Task output formats etc., before guardrails.