Solving the Imbalanced Data Problem in RAG

Pure Python Weighted Random Sampling for Enterprise RAG and Multi-Agent Systems

In enterprise Retrieval-Augmented Generation (RAG) and Multi-Agent systems, we often face a silent killer of LLM performance: Imbalanced Data.

When building few-shot prompts or selecting context windows from a retrieved corpus, standard random sampling will overwhelmingly favor majority classes. If 99% of your retrieved data belongs to "Class A" and 1% to "Class B", a standard random sampler will almost never select "Class B" for the LLM's context, causing the model to hallucinate or fail on critical edge cases.

While libraries like NumPy offer numpy.random.choice, enterprise environments—especially those deploying to edge devices, strict air-gapped networks, or highly optimized microservices—often demand zero external dependencies.

In this end-to-end guide, we will implement a Pure Python Weighted Random Sampler from scratch (using only the standard library). We will then integrate it into an Enterprise Multi-Agent LangGraph RAG System with persistent memory and state, solving the imbalanced context problem for an IT Incident Routing Agent.

The Real-World Use Case: IT Incident Routing & Resolution

Imagine an enterprise IT helpdesk processing 10,000 tickets a day.

The Problem

9,500 tickets are Password Resets. 400 are Software Installs. 100 are Suspected Ransomware.

The RAG Pipeline

To help the LLM draft a response or route the ticket, we retrieve the top 50 most similar historical tickets from the vector database.

The Failure Mode

If we randomly sample 5 of those 50 tickets to use as few-shot examples in the prompt (to save tokens and maintain focus), we will likely get 5 Password Reset examples. The LLM never sees the Ransomware examples, leading to catastrophic misrouting when a real security incident occurs.

The Solution

We use a Weighted Random Sampler. We assign higher weights to minority classes (e.g., Ransomware gets a weight of 50, Password Reset gets a weight of 1). The sampler guarantees that critical, rare examples are proportionally represented in the LLM's context window.

The Algorithm: Pure Python Weighted Sampling

To sample from an imbalanced dataset without external libraries like NumPy, we use the Cumulative Distribution Function (CDF) + Binary Search approach.

Calculate cumulative weights by transforming the raw weights into a running total.
Generate a random target between 0 and the total weight.
Use Python's built-in bisect module to efficiently locate the selected index.

Note: While Python 3.6+ includes random.choices() in the standard library, implementing the CDF manually demonstrates the underlying algorithmic mechanics and allows for custom optimizations (such as caching the CDF for repeated sampling).

End-to-End Python Implementation

Prerequisites

pip install langgraph langchain-openai langchain-core

Step 1: The Pure Python Weighted Sampler

First, we write the core algorithmic function. It relies only on Python's built-in random and bisect modules.

import random
import bisect
from typing import List, Any, Union

def weighted_random_sample(
    population: List[Any],
    weights: List[Union[int, float]],
    k: int = 1
) -> List[Any]:
    """
    Selects k items from a population based on relative weights.
    Implemented from scratch using CDF and Binary Search (No external libs).

    :param population: List of items to sample from.
    :param weights: List of weights corresponding to the population.
    :param k: Number of items to sample.
    :return: List of sampled items.
    """
    if len(population) != len(weights):
        raise ValueError("Population and weights must have the same length.")
    if k > len(population):
        raise ValueError("Sample size k cannot exceed population size.")
    if any(w < 0 for w in weights):
        raise ValueError("Weights must be non-negative.")

    # 1. Build the Cumulative Distribution Function (CDF)
    cum_weights = []
    current_sum = 0.0

    for w in weights:
        current_sum += w
        cum_weights.append(current_sum)

    total_weight = cum_weights[-1]

    if total_weight == 0:
        raise ValueError("Total weight cannot be zero.")

    # 2. Sample k items
    sampled_items = []

    for _ in range(k):
        # Generate a random target between 0 and total_weight
        target = random.uniform(0, total_weight)

        # 3. Binary search to find the index
        index = bisect.bisect_right(cum_weights, target)

        # Edge case: if target is exactly total_weight
        if index >= len(population):
            index = len(population) - 1

        sampled_items.append(population[index])

    return sampled_items

Step 2: Define the LangGraph State and Memory

We define the state to track the incoming ticket, the retrieved historical tickets, the calculated weights, and the final sampled context.

from typing import TypedDict, List, Dict, Any
from langgraph.checkpoint.memory import MemorySaver

class TicketState(TypedDict):
    ticket_id: str
    incoming_text: str

    # Retrieved from Vector DB (Imbalanced)
    retrieved_tickets: List[Dict[str, Any]]

    # Calculated weights for the sampler
    weights: List[float]

    # The final sampled few-shot examples
    sampled_context: List[Dict[str, Any]]

    final_routing_decision: str

Step 3: Build the Multi-Agent Nodes

Node 1: Retriever (Simulated)

Simulates fetching 50 historical tickets where 90% are Password Reset and 10% are Ransomware.

def retriever_node(state: TicketState) -> Dict[str, Any]:
    """Simulates a Vector DB returning an imbalanced set of historical tickets."""
    retrieved = []

    # 45 Password Resets
    for i in range(45):
        retrieved.append({
            "id": f"pwd_{i}",
            "category": "Password Reset",
            "text": "Reset my password please"
        })

    # 5 Ransomware alerts (Rare but critical)
    for i in range(5):
        retrieved.append({
            "id": f"ransom_{i}",
            "category": "Ransomware",
            "text": "Files are encrypted, send bitcoin"
        })

    return {"retrieved_tickets": retrieved}

Node 2: Context Sampler (The Core Logic)

This node calculates the inverse-frequency weights to upsample the rare Ransomware class, then calls our pure Python sampler.

from collections import Counter

def context_sampler_node(state: TicketState) -> Dict[str, Any]:
    """
    Calculates weights based on inverse category frequency and
    uses the pure Python weighted sampler to select diverse few-shot examples.
    """
    retrieved = state["retrieved_tickets"]

    # Count frequencies of each category
    category_counts = Counter([t["category"] for t in retrieved])

    total_tickets = len(retrieved)
    num_categories = len(category_counts)

    # Calculate inverse-frequency weights
    weights = []

    for ticket in retrieved:
        cat_count = category_counts[ticket["category"]]
        weight = total_tickets / (num_categories * cat_count)
        weights.append(weight)

    sampled_tickets = weighted_random_sample(
        population=retrieved,
        weights=weights,
        k=5
    )

    return {
        "weights": weights,
        "sampled_context": sampled_tickets
    }

Node 3: Synthesizer Agent (LLM)

The LLM uses the sampled context to make a routing decision.

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def synthesizer_node(state: TicketState) -> Dict[str, Any]:
    """LLM Agent uses the sampled few-shot context to route the ticket."""

    context_str = "\n".join([
        f"- Example {i+1} ({t['category']}): {t['text']}"
        for i, t in enumerate(state["sampled_context"])
    ])

    prompt = ChatPromptTemplate.from_template("""
    You are an Enterprise IT Routing Agent.

    INCOMING TICKET:
    {incoming_text}

    FEW-SHOT EXAMPLES (Selected via Weighted Sampling for Diversity):
    {context}

    Based on the incoming ticket and the diverse examples provided,
    classify the ticket into one of the following categories:

    [Password Reset], [Software Install], [Ransomware]

    Output ONLY the category name.
    """)

    chain = prompt | llm

    decision = chain.invoke({
        "incoming_text": state["incoming_text"],
        "context": context_str
    }).content

    return {
        "final_routing_decision": decision.strip()
    }

Step 4: Compile the Graph with Memory

We wire the nodes together and attach a MemorySaver. This ensures that if the user follows up (e.g., "Are you sure it's not a software install?"), the agent remembers the previous state, retrieved chunks, and routing decision.

from langgraph.graph import StateGraph, START, END

def build_it_routing_graph():
    workflow = StateGraph(TicketState)

    # Add Nodes
    workflow.add_node("retrieve", retriever_node)
    workflow.add_node("sample", context_sampler_node)
    workflow.add_node("synthesize", synthesizer_node)

    # Define Edges
    workflow.add_edge(START, "retrieve")
    workflow.add_edge("retrieve", "sample")
    workflow.add_edge("sample", "synthesize")
    workflow.add_edge("synthesize", END)

    # Compile with Memory
    memory = MemorySaver()
    return workflow.compile(checkpointer=memory)

app = build_it_routing_graph()

Step 5: Execution and Verification

Let's run the graph with a critical Ransomware ticket. Because of our weighted sampler, the LLM will actually see the rare ransomware examples in its context, rather than being blinded by 45 password resets.

def run_pipeline():
    thread_id = "incident_9942"
    config = {"configurable": {"thread_id": thread_id}}

    # Simulate a critical incoming ticket
    initial_state = {
        "ticket_id": "INC-9942",
        "incoming_text": "URGENT: All my documents have .locked extensions and there is a readme.txt demanding crypto.",
        "retrieved_tickets": [],
        "weights": [],
        "sampled_context": [],
        "final_routing_decision": ""
    }

    print("🚀 Running Multi-Agent RAG Pipeline with Pure Python Weighted Sampling...\n")

    # Stream the execution to see state changes
    for event in app.stream(initial_state, config=config):
        for node_name, state_update in event.items():
            print(f"--- Node: {node_name} ---")

            if node_name == "sample":
                categories_picked = [
                    t["category"]
                    for t in state_update["sampled_context"]
                ]

                print(f"Sampled Categories: {categories_picked}")
                print("(Notice how Ransomware is heavily represented despite being only 10% of the retrieved data!)\n")

            elif node_name == "synthesize":
                print(f"Final Routing Decision: {state_update['final_routing_decision']}\n")

if __name__ == "__main__":
    run_pipeline()

Enterprise Benefits of This Approach

Zero External Dependencies for Core Math

By implementing the CDF + Binary Search algorithm using only Python's random and bisect standard libraries, we eliminate the need for NumPy or SciPy in the sampling layer. This drastically reduces the Docker image size, speeds up cold starts in serverless and edge environments, and removes supply-chain vulnerabilities.

Deterministic Context Window Management

LLM context windows are expensive. Standard random sampling wastes tokens on redundant majority-class examples. Weighted sampling guarantees token efficiency by forcing diversity into the prompt.

Stateful Multi-Agent Orchestration

By wrapping this logic in LangGraph, the sampling process isn't a black box. The TicketState explicitly tracks the retrieved_tickets, the calculated weights, and the final sampled_context. This provides complete observability for enterprise auditing.

Persistent Memory

Because we utilized LangGraph's Checkpointer, the system remembers the exact context it sampled for a specific thread_id. If a human agent needs to review the AI's decision later, they can query the graph state and see exactly which few-shot examples influenced the LLM's routing decision.

Conclusion

Handling imbalanced data in AI systems doesn't always require heavy machine learning libraries. By combining a fundamental algorithmic approach (Pure Python Weighted Sampling via CDF and Binary Search) with modern orchestration frameworks (LangGraph Multi-Agent RAG), we can build highly resilient, context-aware enterprise systems.

This architecture ensures that your LLMs are never blinded by the majority class, critical edge cases are always represented in the context window, and the entire decision-making process remains stateful, memory-backed, and fully auditable.