loke.dev
Header image for A 5-Point Checklist for Hardening Your AI Tool-Calling Logic

A 5-Point Checklist for Hardening Your AI Tool-Calling Logic

Build more resilient LLM integrations by moving beyond basic prompting to a robust architecture of schema-driven validation and state-aware retries for your AI tools.

· 4 min read

The prompt "return a valid JSON object" is the single most ignored instruction in the history of human-computer interaction. You can ask an LLM to follow a schema a thousand times, but eventually, it will hallucinate a field, forget a closing brace, or decide that a date should be "Tuesday-ish" instead of an ISO string.

Building a tool-calling integration that actually works in production isn't about writing the perfect system prompt; it's about building a defensive perimeter around your code. If you're just passing raw LLM output directly into your internal functions, you’re basically giving a caffeinated intern the keys to your production database.

Here is a 5-point checklist to harden your AI tool-calling logic before it breaks something expensive.

---

1. Schema Validation is Your Only Friend

Don't just define your tools in the tools array and hope for the best. Use a library like Pydantic (Python) or Zod (TypeScript) to enforce strict types the second the model returns a call.

The LLM will often try to pass a string when you expect an integer, or it’ll get creative with enum values. By using a validator, you catch these errors before they hit your business logic.

from pydantic import BaseModel, Field, ValidationError
from typing import Literal

class GetWeather(BaseModel):
    location: str = Field(..., description="The city and state, e.g. San Francisco, CA")
    unit: Literal["celsius", "fahrenheit"] = "celsius"

# When the LLM returns arguments:
raw_args = '{"location": "London", "unit": "kelvin"}' # Oops, 'kelvin' isn't allowed

try:
    tool_call = GetWeather.model_validate_json(raw_args)
except ValidationError as e:
    print(f"Caught a hallucination: {e}")

If it fails validation, you don't just crash. You feed that error back to the model (more on that in a second).

2. Implement the "Self-Correction" Loop

When a tool call fails—whether because of a validation error or a failed API request—the most resilient move is to tell the LLM exactly why it failed and ask it to try again.

I've seen so many developers just log an error and return a "Sorry, I can't do that" to the user. Instead, treat the error as a new turn in the conversation.

The workflow looks like this:
1. LLM sends a bad tool call.
2. Your code catches the ValidationError or 400 Bad Request.
3. You append a tool message to the history: "Error: 'unit' must be one of ['celsius', 'fahrenheit']. You provided 'kelvin'."
4. You call the LLM again.

Nine times out of ten, the model sees the mistake and fixes the JSON in the second attempt.

3. Guard Against the "Infinite Loop of Death"

While self-correction is great, you need a circuit breaker. If the model keeps hallucinating the same wrong parameter, it will burn through your API credits faster than a crypto miner.

Always wrap your tool-processing logic in a max_iterations counter.

max_retries = 3
attempts = 0

while attempts < max_retries:
    response = call_llm(messages)
    tool_calls = response.choices[0].message.tool_calls
    
    if not tool_calls:
        break
        
    results = execute_tools(tool_calls)
    
    # If all tools succeeded, we're done. 
    # If there were errors, they are now in 'messages' for the next loop.
    if "error" not in str(results): 
        break
        
    attempts += 1

If you hit attempts == 3, give up and tell the user, "I'm having trouble connecting to that service right now." It’s much better than a $50 bill for a loop that didn't accomplish anything.

4. State-Aware Argument Sanitization

LLMs are notoriously bad at remembering exactly what IDs look like. If you have a tool called update_record(record_id: str), the model might try to pass the name of the record or a placeholder like "record_123".

I learned the hard way that you should always perform a "lookup or fuzzy match" step if the LLM provides an identifier. If the model says "Delete the email from Bob," and your tool expects a UUID, don't just fail. Search for "Bob" in the context, find the UUID, and use that.

Better yet, provide the IDs in the context explicitly so the model doesn't have to guess:

"Available Contacts: Bob (ID: u_882), Alice (ID: u_991)"

5. The "Wait, Are You Sure?" Threshold

Some tools are "Read-Only" (searching Google, checking weather) and some are "Write-Heavy" (deleting a repo, sending an invoice, emailing a boss).

Never let an AI execute a Write-Heavy tool without a manual confirmation or a strict sanity check.

If my AI agent is calling send_payment(), I don't just run it. I have the tool return a "pending" status and show a UI button to the user. Only when the human clicks "Confirm" does the actual logic execute.

If you're building a CLI tool, it looks like this:

def delete_file_tool(filename: str):
    # A simple safety check
    if filename.endswith(".env") or filename == "/":
        return "Error: System files cannot be deleted."
    
    confirm = input(f"AI wants to delete {filename}. Allow? (y/n): ")
    if confirm.lower() == 'y':
        # logic to delete
        return "File deleted."
    return "Action cancelled by user."

It might feel less "magic," but it's the difference between a helpful assistant and a catastrophic bug.

Final Thoughts

Hardening your AI logic is mostly about pessimism. Assume the model will get the JSON wrong. Assume the API will be down. Assume the model will try to delete root.

If you build your architecture around validation, retries, and safety rails, you move from a "cool demo" to a production-grade tool that people can actually trust. Focus on the plumbing, and the intelligence will handle the rest.