Build a Streaming “JSON Repair” Generator for RAG Tool Calls

I ran into a frustrating failure mode while building an app that uses Retrieval-Augmented Generation (RAG)—where we fetch relevant text from a knowledge base and then ask a language model to produce an answer using that text. In my case, the model wasn’t just “sometimes wrong.” It was almost right: it would output JSON for a “tool call” (like search, create_ticket, etc.), but the JSON would be broken in tiny ways—an extra trailing comma, a missing quote, or a truncated string. That broke downstream code hard.

So I built a small, practical system: a streaming JSON repair generator that fixes common structural issues and verifies the result before executing tool calls. Here’s what I learned by implementing it end-to-end.

---

The niche problem: streaming tool-call JSON that breaks by one character

The model’s output looked like this (intentionally malformed):

{
  "tool": "search",
  "query": "hypertension guidelines",
  "top_k": 5,
}

That trailing comma after 5 makes strict JSON parsers throw an error. In real streams, it gets worse: sometimes the JSON arrives truncated mid-string, like:

{
  "tool": "search",
  "query": "hypertension guidelines,
  "top_k": 5
}

My goal wasn’t “make the model produce perfect JSON.” I wanted a resilient layer that:

1. Repairs common syntax issues
2. Works with streaming text (chunks arriving incrementally)
3. Validates the repaired JSON matches a schema I expect
4. Only then allows the tool call to execute

---

What I built: a two-stage pipeline

I used two small components:

1. A streaming accumulator that collects text chunks into a buffer.
2. A “JSON repair” function that:
   • trims to the first complete-looking JSON object
   • removes a few known JSON-breaking patterns
   • tries parsing
   • validates required fields and types

This is not a full general-purpose “JSON hallucination fixer.” It’s a pragmatic repair layer tailored to tool-call JSON in RAG setups.

---

Working code: streaming accumulator + JSON repair + schema validation

1) Install dependencies

pip install jsonschema

2) The complete script

import json
import re
from typing import Dict, Any, Optional, List
from jsonschema import validate


# --- Schema we expect for a tool call ---
TOOL_CALL_SCHEMA = {
    "type": "object",
    "properties": {
        "tool": {"type": "string", "enum": ["search", "create_ticket"]},
        "query": {"type": "string"},
        "top_k": {"type": "integer", "minimum": 1, "maximum": 50},
        "ticket_title": {"type": "string"},
    },
    "required": ["tool"],
    "additionalProperties": True,
}


def extract_first_json_object(text: str) -> str:
    """
    Find the first JSON object substring by locating the first '{'
    and then scanning for a balanced brace sequence.
    """
    start = text.find("{")
    if start == -1:
        raise ValueError("No JSON object start '{' found.")

    depth = 0
    in_string = False
    escape = False

    for i in range(start, len(text)):
        ch = text[i]

        if in_string:
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == '"':
                in_string = False
            continue

        if ch == '"':
            in_string = True
            continue

        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[start:i + 1]

    # If we get here, the buffer does not contain a complete object yet.
    raise ValueError("No complete JSON object yet (brace balance not closed).")


def repair_common_json_issues(s: str) -> str:
    """
    Fix a few common LLM JSON mistakes:
    - trailing commas before } or ]
    - missing quotes around keys are hard; we won't guess blindly
    - common truncation issue is handled by waiting for completion (extract_first_json_object)
    """
    # Remove trailing commas like: { "a": 1, }
    s = re.sub(r",(\s*[}\]])", r"\1", s)

    # Sometimes models emit smart quotes; normalize them.
    s = s.replace("“", '"').replace("”", '"').replace("’", "'")

    return s


def parse_and_validate_tool_call(text: str) -> Dict[str, Any]:
    """
    Turn the repaired JSON text into a Python dict, then validate schema.
    """
    repaired = repair_common_json_issues(text)
    obj = json.loads(repaired)  # strict parsing
    validate(instance=obj, schema=TOOL_CALL_SCHEMA)
    return obj


class StreamingJSONRepairGenerator:
    """
    A minimal “streaming” processor:
    feed() receives text chunks; it tries to extract and parse the first JSON object.
    """
    def __init__(self):
        self.buffer = ""

    def feed(self, chunk: str) -> Optional[Dict[str, Any]]:
        self.buffer += chunk

        # Only attempt parsing once we likely have a complete object.
        # extract_first_json_object will raise until braces are balanced.
        try:
            json_text = extract_first_json_object(self.buffer)
            tool_call = parse_and_validate_tool_call(json_text)
            return tool_call
        except ValueError:
            # Not enough text yet to complete an object
            return None
        except json.JSONDecodeError:
            # Something still malformed; in a production system you could
            # run a more aggressive repair strategy here.
            return None
        except Exception:
            # Schema validation failed or other issues
            return None


if __name__ == "__main__":
    # Simulate a streaming response that includes a trailing comma error.
    chunks: List[str] = [
        '{\n  "tool": "search",\n  "query": "hypertension guidelines",\n  "top_k": 5,',
        "\n}\nExtra trailing text not part of JSON."
    ]

    stream = StreamingJSONRepairGenerator()

    result = None
    for i, c in enumerate(chunks, 1):
        result = stream.feed(c)
        print(f"[chunk {i}] fed {len(c)} chars -> result ready: {result is not None}")

        if result is not None:
            break

    if result is None:
        raise SystemExit("Failed to repair and parse tool call from stream.")

    print("\nParsed tool call:")
    print(json.dumps(result, indent=2))

    # Example “execution” step:
    tool = result["tool"]
    if tool == "search":
        print(f"\nExecuting search with query='{result.get('query')}' top_k={result.get('top_k')}")

---

How the code behaves (step-by-step)

1) Accumulate streaming chunks

The StreamingJSONRepairGenerator maintains a self.buffer. Each time feed(chunk) is called, it appends the chunk.

This matters because the broken JSON may not be parsable until the end arrives.

2) Extract the first complete JSON object

extract_first_json_object() scans from the first { and tracks brace depth, while correctly ignoring braces inside quoted strings.

• If braces aren’t balanced yet, it raises ValueError("No complete JSON object yet...")
• The streaming controller catches that and returns None

This prevents premature parsing attempts on partial JSON.

3) Repair known “LLM JSON” syntax issues

repair_common_json_issues() focuses on a common failure:

• trailing commas before } or ]

That single regex turns:

{ "top_k": 5, }

into:

{ "top_k": 5 }

4) Parse strictly and validate schema

Once repaired, it uses json.loads() (strict parsing) and then jsonschema.validate() against TOOL_CALL_SCHEMA.

This is where you prevent nonsense tool calls like:

{ "tool": "drop_database" }

because schema restricts tool to known enum values.

5) Only execute after validation

In __main__, I show a simplified “execution” step for search. The important rule is: execute only when result is not None.

---

Why this beats “try/catch and pray” in tool-call pipelines

I used to rely on a blanket try: json.loads(...) except: ... strategy. The problem was that it often failed silently or attempted to execute broken tool calls.

This pipeline is more deterministic:

• It won’t parse until braces are balanced (so truncation doesn’t become “sometimes valid garbage”).
• It repairs only known syntax patterns (so it doesn’t invent structure).
• It validates schema so only tool calls that match expected shape are executable.

---

Where RAG fits in

In a real RAG workflow, the model might generate the tool-call JSON based on retrieved documents. The generation can include subtle formatting errors due to prompt formatting, long contexts, or chunk truncation. My repair layer sits between:

retrieved context → model output (stream) → repair + validate → tool execution

That separation is the difference between a fragile demo and something that survives messy real-world model output.

---

Conclusion

I built a streaming JSON repair layer specifically for RAG tool-call outputs, focusing on extracting the first complete JSON object, fixing a small set of common syntax issues (like trailing commas), and enforcing a strict JSON schema before executing tools. The big win for me was reliability: instead of letting broken JSON crash (or worse, silently execute the wrong thing), the pipeline deterministically waits, repairs what it can safely repair, validates shape, and only then proceeds.