DEV Community

Cover image for Towards the Stability of LLM Agents
kai wen ng
kai wen ng

Posted on

Towards the Stability of LLM Agents

An industry-level LLM agent is not simply an API call that returns a response. It needs to be resilient to transient failures, malformed outputs, and schema violations.
To improve the stability of my agent system, I introduced two decorators around my LLM calls. They handle three common failure modes:

  • JSON decoding errors — provide the model with its previous malformed response and ask it to regenerate valid JSON.
  • Pydantic BaseModel validation errors — provide the model with the specific schema validation errors and ask it to correct the response.
  • 429 rate-limit errors — retry with appropriate backoff.

This separates the reliability logic from the individual agent functions, allowing each LLM call to focus on its own task while the decorators provide a consistent recovery mechanism.

Json and Pydantic Basemodel

To handle problematic model outputs, I first use the json-repair package:
json-repair
PS: If you find the package useful, consider giving the author a ⭐ to show some support.

This allows the system to recover from common JSON formatting issues in the model's response. If the repaired result is still empty, I raise a JSONDecodeError manually.

This is useful because the retry decorator can then intercept the exception and provide the model with a meaningful retry message. Instead of blindly asking the model to try again, I can tell it:

Your previous response was not valid JSON.

and include the previous response through e.doc. This gives the model its failed output as context and allows it to regenerate the response correctly.

The next case is Pydantic ValidationError. A response can be valid JSON but still fail to match the expected schema. In this case, simply asking the model to retry is also not particularly useful. The retry message should contain the specific fields that failed validation and the expected constraints.

This information is available through e.errors(). I can extract the field location and validation message, then feed those details back to the model. This turns the retry from a blind regeneration into a targeted correction.


def llm_retry(max_retries=3):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            retry_message = None

            for attempt in range(max_retries):
                try:
                    return await func(*args, retry_message=retry_message, **kwargs)

                except ValidationError as e:
                    errors = [
                        {"field": ".".join(map(str, x["loc"])), "message": x["msg"]}
                        for x in e.errors()
                    ]
                    retry_message = (
                        "Your previous response failed schema validation.\n\n"
                        f"Validation errors:\n{errors}\n\n"
                        "Correct these errors and return the complete response as JSON only."
                    )

                except json.JSONDecodeError as e:
                    retry_message = (
                        "Your previous response was not valid JSON.\n\n"
                        f"Previous response:\n{e.doc}\n\n"
                        "Regenerate the complete response as valid JSON only. "
                        "Do not include Markdown or explanations."
                    )

            raise RuntimeError(f"LLM failed after {max_retries} attempts")

        return wrapper
    return decorator
Enter fullscreen mode Exit fullscreen mode

The 429 Rate-Limit Decorator

This part is relatively straightforward. I use a simple function factory to create a decorator that wraps around the LLM calls and automatically retries requests that receive a 429 Too Many Requests response.

The decorator applies an exponential backoff between attempts, allowing the system to recover from temporary rate limiting without requiring each LLM call to implement its own retry logic.

def Retry429(trials=3, delay=1.0, backoff=2.0):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            current_delay = delay

            for attempt in range(trials):
                try:
                    return await func(*args, **kwargs)
                except RateLimitError:
                    pass
                except httpx.HTTPStatusError as e:
                    if e.response.status_code != 429:
                        raise
                except httpx.HTTPError:
                    raise

                if attempt == trials - 1:
                    raise

                await asyncio.sleep(current_delay)
                current_delay *= backoff

        return wrapper

    return decorator
Enter fullscreen mode Exit fullscreen mode

After having these decorators in place, I can define an LLM call as a simple function. Each function is responsible only for constructing the input arguments and returning the output arguments, represented by LLMRequest and LLMResponse.

Defining these interfaces as Pydantic BaseModels provides a consistent contract between different parts of the system and makes the overall agentic architecture more systematic and maintainable.

More importantly, the retry and recovery mechanisms are no longer coupled to individual LLM implementations. This allows me to support different types of LLM calls, such as image generation and VLM inference, without copy-pasting the same retry and validation logic across every function. The individual functions remain focused on their specific task, while the shared decorators handle the common reliability concerns.


    @Retry429(trials=10, delay=10)
    @llm_retry()
    async def generate(self, llm_request: LLMRequest, retry_message=None)-> LLMResponse:
        args = {}
        # Just some args parsing
        if retry_message is not None:
            args["messages"].append({"role": "system", "content": retry_message})

        response = await self.client.chat.completions.create(**args)
        self.add_token_usage(response)
        if not response:
            raise RuntimeError("LLM returned empty response")

        if not response.choices:
            logger.error(
                "LLM returned no choices. Response=%s",
                response.model_dump()
            )
            raise RuntimeError("LLM returned no choices")
        content = response.choices[0].message.content

        json_schema = None
        if llm_request.json_schema is not None:
            content_json = json_repair.loads(content, return_objects=True) 
            if len(content_json) == 0:
                raise json.JSONDecodeError("Invalid json", content, 0)
            if isinstance(content_json, list):
                json_schema = []
                for content in content_json:
                    json_schema.append(llm_request.json_schema(**content))
            elif llm_request.bulk:
                json_schema = [llm_request.json_schema(**content_json)]
            else:
                json_schema = llm_request.json_schema(**content_json)
Enter fullscreen mode Exit fullscreen mode

I never expected the concepts from Fluent Python—especially data models, function factories, and decorators—to be this useful when designing an elegant system like this.

Seeing how these concepts can be combined to build a systematic and maintainable LLM infrastructure has been a good reminder that strong software engineering fundamentals remain highly relevant when building modern AI systems.

It also encourages me to keep improving my engineering skills by reading more deeply, understanding the fundamentals, and applying what I learn through real projects.

Top comments (0)