Posted by Shubham Saha, AI Researcher of Gaudiy Inc.

Introduction

Probabilistic programming has emerged as a powerful paradigm for building models that reason under uncertainty. In traditional probabilistic ML, practitioners typically handcraft inference algorithms — a laborious and error-prone process that requires deep mathematical expertise. Probabilistic programming changes this: rather than hand-coding complex inference algorithms, practitioners can express their models as programs and let the system handle the mathematical heavy lifting. But here’s the catch: general-purpose inference is often too slow, while hand-tuned algorithms are tedious to implement and maintain.

Enter Gen — a probabilistic programming system developed by the MIT Probabilistic Computing Project that gives you the best of both worlds: expressive modeling with programmable inference.

What is Gen?

Gen is an open-source probabilistic programming system that automates the tricky math and low-level implementation details of inference algorithms while giving users the flexibility to customize solutions for their specific models.

Originally built in Julia (as Gen.jl), Gen lets you write generative models using ordinary code — including stochastic branching, loops, and recursion — while providing building blocks for constructing efficient inference algorithms.

Key capabilities:

• Programmable inference: Compose inference algorithms (MCMC, SMC, variational inference) tailored to your model’s structure — mix Gibbs sampling for discrete variables with HMC for continuous ones
• Hybrid inference: Combine neural network proposals with principled Bayesian inference (out of scope for this post, but a key capability for scaling to complex domains)
• Stochastic structure: Model problems where the number of latent variables is unknown (e.g., object detection, clustering with unknown K)
• Programmable trade-offs: Use high-level abstractions for prototyping, drop down to optimized code for production
• Clean API: Automatic differentiation and inference primitives

Enter GenJAX: GPU-Accelerated Bayesian Inference for Python

GenJAX reimplements Gen’s core abstractions directly on JAX, bringing principled uncertainty quantification to Python with the performance characteristics production AI demands. For practitioners already working in the Python ecosystem, this means programmable inference without leaving your stack — leveraging JAX’s strengths:

• GPU acceleration: JIT compilation and automatic batching of inference computations onto GPU devices
• JAX transformations: Full compatibility with jit, vmap, and automatic differentiation
• Familiar syntax: Use the @gen decorator to define generative functions in Python
• Built-in algorithms: Sampling Importance Resampling (SIR) and other inference methods out of the box

A Quick Taste

In GenJAX, you define generative functions — computational objects representing probability measures over structured sample spaces. These functions can be composed, traced, and used with various inference algorithms:

@gen
def my_model():
    slope = normal(0.0, 1.0) @ "slope"
    intercept = normal(0.0, 1.0) @ "intercept"
    # ... rest of the model

Why This Matters

Traditional probabilistic programming systems offer limited control over inference: you either accept slow, general-purpose algorithms or implement everything from scratch. Gen’s programmable inference approach lets you compose algorithm components — importance sampling, MCMC kernels, gradient-based methods — to build inference strategies matched to your model’s structure.

This is particularly valuable for:

• Bayesian nonparametric models with unbounded parameter spaces
• Structured prediction where the number of latent variables is itself unknown
• Scientific modeling requiring interpretable uncertainty quantification
• Robotics and perception systems that must reason about 3D scenes

Hands-On Example: Bayesian House Price Prediction

Let’s put GenJAX to work on a real problem: predicting house prices using the Kaggle House Prices dataset. This dataset contains 79 features describing residential homes in Ames, Iowa, with the goal of predicting SalePrice.

Instead of a point estimate from traditional regression, we’ll build a Bayesian model that gives us uncertainty quantification — knowing not just what the predicted price is, but how confident we should be in that prediction.

The Problem Setup

The Kaggle dataset includes features like:

• GrLivArea: Above-ground living area (square feet)
• OverallQual: Overall material and finish quality (1-10)
• YearBuilt: Original construction date
• TotalBsmtSF: Total basement square footage

For simplicity, we’ll focus on a subset of numerical features and build a Bayesian linear regression model.

Defining the Generative Model

In GenJAX, we express our beliefs about how house prices are generated:

import jax
import jax.numpy as jnp
import jax.random as jrandom
from genjax import gen, normal, uniform, Target, ChoiceMap
from genjax.inference.smc import ImportanceK

@gen
def house_price_model(X):
    """
    Bayesian linear regression for house prices.
    X: feature matrix (n_samples, n_features)
    """
    # Priors on regression coefficients
    # We expect positive coefficients for quality/size features
    coef_0 = normal(0.0, 1.0) @ "coef_0"
    coef_1 = normal(0.0, 1.0) @ "coef_1"
    coef_2 = normal(0.0, 1.0) @ "coef_2"
    coef_3 = normal(0.0, 1.0) @ "coef_3"
    # Prior on intercept (log-scale, since we'll predict log prices)
    intercept = normal(12.0, 1.0) @ "intercept"  # ~$160k baseline
    # Prior on noise (houses vary in price even with same features)
    noise_std = uniform(0.1, 0.5) @ "noise_std"
    # Generate predictions for each house
    coeffs = jnp.array([coef_0, coef_1, coef_2, coef_3])
    predictions = X @ coeffs + intercept
    # Likelihood: observed prices given our linear model
    # GenJAX supports vectorized distributions - the normal distribution
    # broadcasts over the predictions array, treating each as independent
    log_prices = normal(predictions, noise_std) @ "log_prices"
    return predictions

This model encodes several assumptions:

1. Coefficients are uncertain: We don’t know the true relationship, so we place priors
2. Prices are log-normal: Prices can’t be negative and are right-skewed (we model log prices, making multiplicative effects additive)
3. Heterogeneous noise: Even similar houses sell for different prices

Note: The model uses explicit coefficient variables and vectorized likelihood for JAX compatibility. The vectorized normal(predictions, noise_std) broadcasts over all samples efficiently. See the full example code for the complete implementation.

Running Inference

Given observed sale prices, we want to infer the posterior distribution over coefficients:

import pandas as pd
import numpy as np

# Load and preprocess data
train_df = pd.read_csv("train.csv")
features = ["GrLivArea", "OverallQual", "YearBuilt", "TotalBsmtSF"]
X = train_df[features].fillna(0).values.astype(np.float32)
y = np.log(train_df["SalePrice"].values).astype(np.float32)
# Standardize features
X_mean, X_std = X.mean(axis=0), X.std(axis=0)
X = (X - X_mean) / X_std
# Convert to JAX arrays
X_train = jnp.array(X)
y_train = y
# Build constraint map with observed prices (vectorized)
constraints = ChoiceMap.d({"log_prices": y_train})
# Create inference target
target = Target(
    house_price_model,
    (X_train,),
    constraints
)
# Run importance sampling with multiple particles
key = jrandom.PRNGKey(42)
k_particles = 1000
alg = ImportanceK(target, k_particles=k_particles)
sub_keys = jrandom.split(key, k_particles)
_, posterior_samples = jax.vmap(alg.random_weighted, in_axes=(0, None))(sub_keys, target)

What We Get: Uncertainty, Not Just Predictions

The beauty of this approach is that we don’t get a single coefficient estimate — we get a distribution:

# Extract posterior samples for coefficients
coef_samples = posterior_samples["coef_0"]  # GrLivArea coefficient

Running the full example on the Kaggle dataset with 1,000 importance sampling particles (using a 70/30 train/test split with 1,021 training samples and 439 test samples) gives us:

Posterior estimates for coefficients:
--------------------------------------------------
    GrLivArea      :   0.116 ± 0.247  (95% CI: [-0.346, 0.588])
    OverallQual    :   0.174 ± 0.284  (95% CI: [-0.395, 0.729])
    YearBuilt      :   0.091 ± 0.236  (95% CI: [-0.379, 0.542])
    TotalBsmtSF    :   0.043 ± 0.240  (95% CI: [-0.395, 0.513])
    Intercept      :  12.026 ± 0.197
    Noise Std      :   0.398 ± 0.075

Holdout predictions with uncertainty (test set):
--------------------------------------------------
    Showing 10 of 439 holdout predictions:
    House | Actual Price  | Predicted (Mean) | 90% Credible Interval
    -----------------------------------------------------------------
        1 | $    112,000 | $        97,850 | $    34,101 - $   271,070
        2 | $    123,600 | $       121,401 | $    41,657 - $   355,248
        3 | $    240,000 | $       161,917 | $    47,347 - $   497,824
        4 | $    117,000 | $       116,189 | $    51,273 - $   263,998
        5 | $    281,213 | $       291,776 | $   105,307 - $   835,386
        6 | $    275,500 | $       266,397 | $    99,238 - $   685,142
        7 | $    130,000 | $       122,954 | $    54,602 - $   296,081
        8 | $    345,000 | $       322,223 | $   117,288 - $   957,747
        9 | $    123,000 | $       125,042 | $    54,339 - $   300,636
       10 | $    179,200 | $       187,161 | $    82,738 - $   448,681
Test set summary metrics (all 439 samples):
--------------------------------------------------
    Mean Absolute Error (log prices): 0.1209
    Mean Absolute Error (actual prices): $21,535
    90% Credible Interval coverage: 99.8%

Notice how the credible intervals (CI) capture the actual prices — even on held-out data the model has never seen. The high coverage on the credible intervals shows the model provides honest uncertainty quantification that’s far more useful for decision-making than a single point estimate.

Why Uncertainty Changes Decisions

Consider House 134 vs House 135 from our results:

A point-estimate model would say both houses are worth ~$350k — virtually identical. But with uncertainty quantification, decision-makers see a completely different picture:

Mortgage lender: For House 135, the wide interval ($97k-$1.3M) signals high valuation uncertainty — the model’s uncertainty reflects how well the observed features explain prices for similar houses in the training data. A lender might require a 25% down payment instead of the standard 20%, or order additional appraisals. For House 134, the tighter interval provides more confidence, allowing standard terms — even though the predicted price is the same.

Home buyer (let’s assume our home buyer is both data savvy and has access to a lot of housing data of his neighborhood): If both houses are listed at $380k (about 9% above prediction), the uncertainty tells different stories. For House 134, with its tighter CI, $380k sits well within the plausible range — reasonable to proceed. For House 135, with its wide CI ($97k-$1.3M), the same $380k asking price is harder to evaluate — the true value could be much lower or the house could actually be a bargain at $380k. Either way, the buyer should investigate further before committing.

Real estate investor: What represents risk to a lender can represent opportunity to an investor. House 135’s wide uncertainty flags it as a property where the model lacks confidence — potentially due to unique characteristics that create arbitrage opportunities. An investor might prioritize investigating this house to find whether it’s mispriced, while House 134’s tighter bounds suggest the market has likely priced it correctly.

The key insight: identical predicted prices with different uncertainty bounds should lead to different actions. Point estimates hide this crucial information.

Extending the Model: Handling Outliers

Real estate data often has outliers. Now imagine a 2,000 sqft house in a good neighborhood that our model predicts should sell for $450,000 — but it actually sold for $180,000 because it was a foreclosure. Without outlier handling, this single data point would drag our coefficients away from their true values, hurting predictions for all normal sales.

In a traditional workflow, you’d need to derive new inference equations from scratch. But with GenJAX, I can simply extend the model with stochastic branching:

from genjax import switch

# Normal houses: tight Gaussian around prediction
@gen
def inlier_obs(mean, noise_std):
    return normal(mean, noise_std) @ "obs"
# Outliers: broad uniform distribution
@gen
def outlier_obs(mean, noise_std):
    return uniform(10.0, 14.0) @ "obs"
# switch selects branch by index: 0 → inlier_obs, 1 → outlier_obs
obs_model = switch(inlier_obs, outlier_obs)
@gen
def robust_house_price_model(X):
    # ... coefficients and intercept as before ...
    for i in range(X.shape[0]):
        # Is this house an outlier?
        is_outlier = flip(0.05) @ f"outlier_{i}"
        # switch requires integer index; flip returns bool
        idx = is_outlier.astype(jnp.int32)
        log_price = obs_model(idx, (predictions[i], noise_std), (predictions[i], noise_std)) @ f"y_{i}"

This is where Gen’s programmable inference shines — you can mix Gibbs sampling for the discrete outlier indicators with Hamiltonian Monte Carlo (HMC) for the continuous coefficients.

Key Takeaways

• Practical applicability: As shown in the house price example, GenJAX can tackle real problems while providing honest uncertainty estimates
• Programmable inference: Gen lets you control the algorithm while automating the math — combine importance sampling with MCMC kernels as needed
• Uncertainty quantification: Instead of point estimates, get full posterior distributions with credible intervals — far more useful for real-world decision-making
• Expressive modeling: Handle complex scenarios like outlier detection through stochastic branching, mixing discrete and continuous inference

Conclusion

Gen represents a shift in how we think about probabilistic programming: instead of treating inference as a black box, it gives practitioners the tools to build custom algorithms tailored to their models. With GenJAX, these capabilities are now accessible to the broader Python/AI community with the performance benefits of JAX.

Whether you’re building Bayesian models for scientific research or developing perception systems that need to reason under uncertainty, GenJAX offers a compelling middle ground between flexibility and automation. At Gaudiy, we focus heavily on AI Agents, and probabilistic programming offers a principled framework for quantifying uncertainty in such systems and hence help improve calibration and robustness.

References

• Gen Official Website
• GenJAX Documentation
• GenJAX GitHub Repository
• Gen.jl GitHub Repository
• MIT Probabilistic Computing Project
• Gen Tutorials
• Probability & Probabilistic Computing Tutorial
• Kaggle House Prices Competition
• Cusumano-Towner, M. F., Saad, F. A., Lew, A., & Mansinghka, V. K. (2019). Gen: A General-Purpose Probabilistic Programming System with Programmable Inference. PLDI 2019.

More Articles：

Gaudiy AI Lab

About Company：

Gaudiy Inc. | Moving forward in a new era with fans

GenJAX: Probabilistic Programming with Programmable Inference was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Sandeep Kumar Nayak, AI Engineer of Gaudiy Inc.

Introduction

Hi, I’m an engineer on the AI team at Gaudiy.

When developing applications with Text-to-Speech (TTS), there’s one problem you’ll inevitably face: proper noun pronunciation issues.

For example, when an English TTS reads the character name “AI” (pronounced “Ah-ee”) from the anime “Oshi no Ko”, it might pronounce it as “A.I.” (artificial intelligence) instead of spelling out “A-I” as two separate letters. For user-facing audio experiences, this is a critical problem that can break immersion and confuse users.

In this article, I’ll provide a comprehensive guide covering:

1. The nature of TTS pronunciation problems — Why they occur and their impact
2. Available solutions — Comparing different approaches
3. Our implementation — A custom pronunciation dictionary for Cartesia TTS
4. Best practices — Lessons learned from production deployment

Whether you’re building a voice assistant, audiobook narrator, or any speech-enabled application, this guide will help you tackle pronunciation challenges effectively.

Background: How TTS Works

Before diving into the problem, let’s understand how modern TTS systems work.

The TTS Pipeline

Modern neural TTS systems typically follow this pipeline:

The critical component for pronunciation is the Grapheme-to-Phoneme (G2P) conversion. This is where the system decides how to pronounce each word. G2P models are trained on dictionaries and text corpora, which means they learn “common” pronunciations but struggle with:

• Out-of-vocabulary words: Names, brands, and neologisms not in training data
• Heteronyms: Words spelled the same but pronounced differently based on context
• Domain-specific terminology: Technical jargon, acronyms, and abbreviations

Why Neural TTS Makes This Harder

Ironically, modern neural TTS systems can make pronunciation problems worse than older concatenative systems. Here’s why:

1. End-to-end learning: Neural models learn pronunciation implicitly, making it harder to override specific words
2. Context-dependent: The model considers surrounding words, which can cause unexpected pronunciation changes
3. Black box behavior: It’s difficult to predict how the model will pronounce a given input

This is why explicit pronunciation control mechanisms are essential.

The Nature of Pronunciation Problems

Why TTS Mispronounces Proper Nouns

TTS engines are trained on massive text datasets, but proper noun pronunciation presents several challenges:

Homographs
Example: AI (person’s name)
Expected: A-I (spelled out)
Actual: “eye” or “Aa-ee”

Acronym ambiguity
Example: GUI
Expected: G-U-I
Actual: “gooey”

Domain-specific terms
Example: NFT
Expected: N-F-T
Actual: “nft” (as one word)

Character names
Example: Naruto
Expected: “Nah-roo-toh”
Actual: “Nah-roo-toe”

The most problematic cases are words that should be pronounced differently based on context.

Real-World Impact

Poor pronunciation has significant consequences:

1. User confusion: Users may not understand what’s being said
2. Brand damage: Mispronouncing your own product name looks unprofessional
3. Accessibility issues: Users relying on TTS for accessibility are most affected
4. Immersion breaking: In entertainment applications, wrong pronunciations destroy the experience

Concrete Example: The “AI” Pronunciation Problem

Let’s examine the “AI” problem in detail. The string “AI” requires different pronunciations depending on context:

Context 1: Artificial Intelligence

Input: "AI is transforming the world"
Expected: "A.I. is transforming the world" (pronounced as initialism)

Context 2: Japanese Name

Input: "AI hoshino is a famous idol"
Expected: "Ai(Ah-ee) Hoshino" is a famous idol" (pronounced as name)

Context 3: Our Application

Input: "AI hoshino is a famous idol"
Problem: TTS says "eye hoshino" or "A.I. hoshino"
We need: "Ai(Ah-ee) Hoshino " (letters spelled out, proper capitalization)

The challenge is that “AI” appears in our training data overwhelmingly as “artificial intelligence”, so the TTS defaults to that pronunciation. We need to override this behavior for our specific use case.

The Substring Problem

Another critical issue is substring conflicts. Consider these dictionary entries:

{
  "AI": "<<a|i>>",
  "AI hoshino": "<<a|i>> Hoshino"
}

If we’re not careful, “AI hoshino” might become “<<a|i>> hoshino” (note the lowercase “h”) because “AI” matches first. This is why pattern ordering matters — a theme we’ll return to throughout this article.

The Solution: Custom Pronunciation Dictionary

1. Phonetic Symbols (IPA)

International Phonetic Alphabet provides precise pronunciation:

AI hoshino → /eɪ aɪ hoʊˈʃiːnoʊ/

Pros:

• Linguistically accurate
• Language-independent

Cons:

• Requires specialized knowledge
• Not all TTS systems support IPA input
• Difficult to maintain for non-linguists

Text Preprocessing (Our Approach)

Replace problematic text before sending to TTS:

text = "AI hoshino is a famous idol"
text = apply_pronunciation_rules(text)
# Result: "<<a|i>> Hoshino is a famous idol"

Pros:

• TTS-agnostic (can switch providers)
• Simple to implement and maintain
• Dynamic updates without restart
• Full control over replacement logic

Cons:

• Requires TTS-specific notation for complex cases
• Must handle pattern conflicts carefully

We chose text preprocessing with TTS-specific notation because it offers the best balance of flexibility, maintainability, and control.

Understanding Cartesia’s Phonetic Notation

Cartesia TTS provides several mechanisms for controlling pronunciation:

1. Sound Splitting with << >>

The most useful feature is splitting words into separate sounds:

"AI" → pronounced as one sound ("eye")
"<<a|i>>" → pronounced as separate letters ("A" then "I")

The | character separates distinct phonetic units. This is perfect for acronyms and initialisms.

2. Emphasis and Pacing

Cartesia also interprets punctuation for pacing:

"Hello, world" → natural pause after comma
"Hello... world" → longer pause with ellipsis
"Hello! World!" → emphatic pronunciation

3. Handling Special Characters

Some characters cause issues and should be removed:

DISALLOWED_CHARACTERS = "「」『』（）[]{}()"

These brackets and parentheses can confuse the TTS model and cause unexpected pauses or pronunciation changes.

Implementation Details

Now let’s dive into the implementation. We’ll build a complete pronunciation override system.

System Architecture

Pronunciation Dictionary Data Structure

The pronunciation dictionary uses a hierarchical JSON structure organized by language:

{
  "ja": {
    "御社": "おんしゃ",
    "弊社": "へいしゃ",
    "漢字": "かんじ",
    "Gaudiy": "ゴーディ"
  },
  "en": {
    "AI hoshino": "<<a|i>> Hoshino",
    "Gaudiy": "Gaudy",
    "NFT": "<<n|f|t>>",
    "GUI": "<<g|u|i>>",
    "API": "<<a|p|i>>"
  },
  "fr": {
    "AI": "a i",
    "Gaudiy": "Godi"
  },
  "es": {
    "AI": "a i"
  },
  "de": {
    "AI": "a i"
  }
}

Design decisions:

1. Language as top-level key: Allows language-specific pronunciations
2. Simple string mapping: Easy to understand and maintain
3. Extensible: Adding new languages requires no code changes
4. Flat structure per language: Simplifies pattern matching

The Critical Rule: Longer Matches First

This is the most important concept in pronunciation dictionary implementation.

The Problem

When you have overlapping patterns, the order of application determines the result:

# Dictionary entries
{
    "AI": "<<a|i>>",
    "AI hoshino": "<<a|i>> Hoshino"
}

Scenario A: Shorter pattern first (WRONG)

Input:  "AI hoshino is a famous idol"
Step 1: Match "AI" → "<<a|i>> hoshino is a famous idol"
Step 2: Match "AI hoshino" → NO MATCH (already replaced!)
Result: "<<a|i>> hoshino is a famous idol"  ❌
        └── lowercase "hoshino" is wrong!

Scenario B: Longer pattern first (CORRECT)

Input:  "AI hoshino is a famous idol"
Step 1: Match "AI hoshino" → "<<a|i>> Hoshino is a famous idol"
Step 2: Match "AI" → no remaining matches
Result: "<<a|i>> Hoshino is a famous idol"  ✅
        └── proper capitalization preserved!

The Solution

Always sort patterns by length (descending) before applying:

sorted_items = sorted(
    replacements.items(),
    key=lambda x: len(x[0]),
    reverse=True  # Longest first
)

This ensures that “AI hoshino” (11 characters) is always checked before “AI” (2 characters).

Future Improvements

Planned Enhancements

1. Context-aware pronunciation: Use NLP to determine pronunciation based on surrounding words
2. User feedback integration: Allow users to report pronunciation issues directly
3. A/B testing: Test pronunciation changes with subset of users
4. Automatic dictionary suggestions: ML model to suggest new entries based on user queries
5. Real-time updates: WebSocket-based dictionary updates without restart

Potential Architecture Evolution

Conclusion

TTS pronunciation issues are a common challenge that can significantly impact user experience. In this article, we’ve covered:

Key Takeaways

1. Understand the problem: TTS models learn “common” pronunciations that may not fit your domain
2. Choose the right approach: Text preprocessing with TTS-specific notation offers the best balance of flexibility and maintainability
3. Implement correctly: The longer match first rule is critical for avoiding substring conflicts

Results

After implementing this system:

• 100% accuracy on known pronunciation issues
• < 1ms overhead for preprocessing
• Zero downtime dictionary updates via GCS
• Multi-language support with language-specific rules

The custom pronunciation dictionary has become an essential component of our TTS pipeline, ensuring that proper nouns and domain-specific terms are pronounced correctly every time.

I hope this comprehensive guide helps you tackle TTS pronunciation challenges in your own applications. If you have questions or want to share your own solutions, feel free to reach out!

Thanks for reading!

References

• Cartesia TTS Documentation
• International Phonetic Alphabet
• Sudachi — Japanese Morphological Analyzer

More Articles：

Gaudiy AI Lab

About Company：

Gaudiy Inc. | Moving forward in a new era with fans

Improving Custom Pronunciation of Voice Agents was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Sandeep Kumar Nayak, AI Engineer of Gaudiy Inc.

Introduction

In recent years, with the growing adoption of LLM-based conversational AI, “contextual understanding of conversation” has become increasingly important. However, improving this ability takes more than new models. True contextual understanding depends on the quality of the dataset that trains and evaluates it. Real conversations are rarely self-contained — they rely on shared history, emotional tone, and subtle references. A strong dataset captures these nuances, enabling models to retrieve the right memories and respond meaningfully. In short, advancing contextual AI starts with data that reflects how people actually communicate. In this article, we introduce our approach to improving memory retrieval through contextual queries at Gaudiy by leveraging a carefully designed dataset. This dataset enables us to quantify performance on contextual memory query tasks and drive meaningful improvements in product quality.

About Gaudiy and Its Product

At Gaudiy, we’re integrating AI-based interactive chat agents into our products to enable meaningful, sustained conversations with users. Our goal is to make these AI-driven interactions feel deeply personal, responsive, and memorable.

The Gaudiy AI team consists of researchers and engineers working across natural language processing, retrieval systems, and human-AI interaction design. Our mission is to build emotionally intelligent AI that adapts to users in real time — leveraging memory and personalization to create rich, dynamic conversations. We combine academic research with product-driven experimentation to bring innovations from lab to life. One of our current focus areas involves experimental conversational features that allow users to engage with AI agents representing specific characters in collaboration with rights holders.

However, user messages are often vague or context-light (e.g., “Nice!” or “Really?”). In such cases, the AI risks retrieving irrelevant or generic content unless it understands the message in the context of recent dialogue. This presents a fundamental challenge in AI conversations: handling ambiguous user inputs accurately.

To tackle this, one of our key areas of improvement is the chat agent’s ability to understand and respond appropriately based on the context of prior conversations. A crucial part of this personalized experience is the AI’s capacity to retain and retrieve long-term memories — shaping ongoing dialogues into richer, more coherent, and contextually relevant interactions.

Why Contextual Queries Matter — and Why the Right Dataset Is Key

Understanding context lies at the heart of delivering meaningful conversational AI experiences. In some setups, such as when the entire prior conversation is stored directly in the prompt, there may be no need for separate memory retrieval steps.

However, to make systems scalable and cost-efficient, most agents rely on short-term context windows. For deeper personalization and long-term coherence, we must store and retrieve relevant long-term memories separately. Identifying which memories are contextually relevant — and retrieving them appropriately — requires an extra step: memory querying. To evaluate how well a system can perform this step, the role of the dataset becomes crucial.

We aim for our dataset to support evaluation across key dimensions that matter for real-world deployment:

1. More appropriate responses: Contextual understanding allows the AI to interpret queries in light of recent conversation history — improving clarity, reducing miscommunication, and increasing user engagement.
2. Memory retrieval accuracy: A well-structured dataset helps benchmark how effectively AI agents retrieve relevant long-term memories.
3. Quantitative Comparison across models: With ground truth and structured dialogue, our dataset enables model benchmarking using rigorous metrics such as MAP@10 .
4. Model selection trade-offs: Our dataset lets us balance performance with cost and latency, guiding embedding model choices based on quantified results.

Illustrative Example: Why Context & Memory Matters

Consider this snippet from a conversation between AI assistant and a user.

User: Have you read any new books recently?
AI: Yes, I recently started reading Harry Potter.
User: That's great. Which volume are you reading?
AI: The first one. It's really interesting!
User: Nice!

This ensures that the AI retrieves relevant details — such as prior mentions of Harry Potter — and can respond with something like: “Yes! I’m excited to continue the story tonight.” rather than a generic affirmation.

Long Term Memory History

Now consider a case where long-term memory comes into play

User: I’m feeling a bit anxious about tomorrow.

Here, the phrase “your big audition” was not mentioned in the current chat. It was something the user mentioned weeks ago:

Long-Term Memory in Database: "User is preparing for a regional theater
audition on June 10th. It's important to them."

Retrieving this information allows the AI to reference it meaningfully, even though it’s outside the recent dialogue window. Without this retrieval, the AI would respond generically. With long-term memory, it might say:

“You’ve worked so hard for this audition — I believe in you.

What kind of dataset helps here?

• It includes multiple turns of chat history.
• It contains annotations or ground-truth for the intended memory reference. ie It simulates long-term memory entries linked to prior user sessions.
• It contains evaluation annotations about which memory entries are contextually relevant given ambiguous prompts.

This showcases the power of memory queries — not just for interpreting immediate context, but for maintaining emotional continuity and personalization over time.

In short, building a high-quality dataset isn’t just about data collection — it’s about unlocking measurable progress in contextual AI by enabling rich, context-aware evaluation

Building on Previous Work

In a prior internal experiment, one of our AI engineers explored contextual query rewriting techniques[1] — including single-shot, multi-shot, and hypothetical reply approaches. This exploration demonstrated how prompt engineering can significantly influence the relevance of memory retrieval, particularly by reducing embedding distances between user queries and relevant memories.

Cosine similarity remains a useful metric for evaluating retrieval effectiveness. This was complemented by qualitative checks — assessing whether the retrieved memory intuitively matched the user query within context. While this method supports rapid iteration and insight into model behavior, it is difficult to ensure consistent, comparable evaluation across different strategies.

To address this, our current work introduces a dataset specifically designed for grounded and measurable evaluation of contextual memory retrieval. Unlike existing public datasets, it targets ambiguous or open-ended Japanese conversational queries, providing ground-truth expectations for what should be retrieved. This allows us to objectively compare retrieval approaches — for example, using metrics such as Mean Average Precision (MAP) — and conduct comprehensive, reproducible evaluations of model performance.

Our Solution

To address these gaps, we developed a private Japanese conversation dataset with the following features:

• Rich conversation history across multiple turns.
• Simulated long-term memory per user.
• Personalization layers including personality traits

This dataset enables quantitative evaluation of how well embedding models perform when handling contextual memory retrieval tasks.

Methodology

1. Data Preparation

We created a context-enriched Japanese dialogue dataset involving:

- Multi-turn conversations.
- Ground-truth memory references.
- Annotated character personas.

Each user query in the dataset is paired with:

- A non-contextual version (standalone user input).
- A contextual version (query + prior conversation history).

This setup lets us measure retrieval quality under real-world, ambiguous input scenarios by directly comparing context-free and context-enriched queries.

2. Performance Metric

To evaluate retrieval, we used:
- Mean Average Precision @10 (MAP@10): Measures how well the top 10 retrieved memories align with ground-truth relevant memories.

Results

1. Contextual Queries Improve Model Performance

Our analysis revealed compelling evidence that our dataset plays a crucial role in demonstrating the need for context-aware memory retrieval. Unlike public datasets, which often lack conversational continuity, our dataset includes rich contextual traces — making it possible to observe the discrepancies in memory retrieval performance when using different embedding models.

Figure 1 : MAP@10 for contextual and non-contextual queries

We found that all models showed a clear improvement in precision when using contextual queries. Furthermore, some embedding models demonstrated significantly greater gains than others. Therefore, in Figure 2, we compare the relative performance improvements to identify which models achieved the highest increase.

Figure 2 : MAP@10 improvement across models of various sizes.

As from both the figures, we see, across all models, adding context to the queries (contextual queries) improved retrieval performance significantly:
Average improvement in MAP: +20%
Top model improvement: Nearly +30%

Insight 1: Even weaker models with poor non-contextual performance showed consistent gains — highlighting how contextual understanding universally boosts performance.

2. Model Selection via Dataset Evaluation

As from the Figure 1, open source models see huge jumps in MAP when context is included — bringing them close to top-tier models like:

• OpenAI/text-embedding-3-large
• Google_gekkko_001

Context Reduces the Dependence on Model Scale:

• The performance advantage of expensive, large-scale models diminishes when contextual understanding is introduced at the query level.
• This implies better ROI for real-world applications using smaller models with proper contextualization.

Insight 2
Use a smaller model
Pre-process user queries with contextual enrichment
Achieve similar accuracy with lower cost and latency

Conclusion

Through our newly introduced context-rich dataset, we have been able to quantitatively measure the performance of different text embedding models in our use case. The results not only validate the effectiveness of context-aware retrieval but also reveal each model’s capability to leverage conversational context — something that public datasets often fail to capture due to their lack of dialogue continuity.

While it has enabled evaluation of memory retrieval, we look ahead to also include the absence of end-to-end output scoring (e.g., user satisfaction) and unmeasured latency-cost tradeoffs — both of which we plan to address.

By enabling quantifiable evaluation on our dataset, we gain a surgical view into memory-retrieval methods and embedding models. Optimizing both in tandem yields systems that produce context-rich, personalized responses without inflating infrastructure cost — raising the bar for conversational AI beyond Gaudiy’s use case.

Looking ahead, we aim to integrate user-level evaluation metrics and investigate lower-cost, high-context embedding models to further improve intelligence and scalability. To ensure these results stand up to scrutiny, we are preparing to submit our findings for peer review. We’ll share updates as that process unfolds — stay tuned as we continue to push the boundaries of context-aware AI.

References

[1] https://medium.com/gaudiy-web3-and-ai-lab/b628896107aa

[2] Mean Average Precision, MAP@10 , https://www.evidentlyai.com/ranking-metrics/mean-average-precision-map

[3] Embedding Models in NLP: https://aws.amazon.com/what-is/embeddings-in-machine-learning/

[4] E5 :https://arxiv.org/abs/2402.05672

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

The Power of Contextual Datasets: Quantifying Memory Retrieval Performance was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Malkoc Ognjen, AI Researcher of Gaudiy Inc.

In the first part of this series on behavior tree agents, we looked at how behavior trees can bring structure to the often unpredictable world of Large Language Models (LLMs), enabling complex and emergent behaviors. Now, we’re building on that foundation to create something practical: a calculator that understands everyday language. Starting from that base, we’ll explore how the inherent prioritization in behavior trees can be harnessed to interpret natural language and guide meaningful computation.

The core idea put simply; when a task can be solved with explicit instructions in a programming language (rather than relying on the semantic interpretation of an LLM) the outcome is almost always more accurate and reliable. This is especially true for mathematics, where precision is critical.

An illustrative example is arithmetic operations in LLMs. This capability was an emergent property that appeared as models grew larger. More recent work has focused on pinpointing the mechanism: distinct activation circuits seem to support operations like addition or multiplication. This this also points to a limitation — while LLMs can manage basic arithmetic, extending that ability to perform more complex, less common, or out‑of‑distribution operations may be increasingly difficult to learn.

Moreover, reinforcement learning that aligns LLMs by rewarding output containing convincing “reasoning” traces does not appear to improve things. The activation circuits responsible for basic operations seem to remain unchanged, or at least disconnected from the reasoning output⁰. In addition, recent studies ² ³ ⁴ ⁵ further suggest that the benefits of reasoning are not universal. Reference 5, in particular, even links performance gains to specific pre‑training. This implies that if a model is not sufficiently large, or has not been exposed to similar tasks during pre‑training, relying on reasoning alone may be ineffective. It is worth exploring alternatives.

With that in mind, let’s do so. Let’s make the LLM assume the role of a stressed-out doctoral student who does not want to, yet again, derive a symbolic expression, or calculate a well-known, but tedious, integral.
Let’s make use of the collective knowledge we already have at our disposal, tools that have been built specifically for this purpose. In particular, the Python programming language offers several advantages: it has a simple syntax, a vast ecosystem of scientific libraries — and, as luck would have it, there’s also an ever‑growing volume of Python code available online — ample material that feeds directly into the pre‑training phase.

With that back‑story, here’s the plan: we’ll use a behavior tree to enforce the core principle introduced earlier. The rule is straightforward: if a task can be solved programmatically, we solve it that way; for everything else, the LLM steps in as a fallback for solving tasks.
Here we will focus on mathematics and physics. When possible, we hand off numerical work to Python, where well‑chosen function can deliver dependable results. The high‑level idea is what we might call a semantic calculator: for tasks the model itself cannot reliably compute, the LLM acts as a translator between the user and the computer.

The handing off work to Python, and resuming translation by an LLM is trivial since we can leverage the prioritization built into behavior trees. This setup lets us route queries intelligently: whenever a task is computable and expressible in Python, we default to that path, since the results are more dependable. Should that route fail, whether during function generation or result retrieval, we intentionally fall back to the LLM’s natural‑language reasoning. The design ensures a baseline accuracy no worse than using the LLM alone, while delivering clear gains whenever a task can be programmatically computed.

The figure above illustrates how user queries are handled. At the root, the system first checks whether the task can be solved programmatically. If so, the query is routed through the Solve Programmatically sequence, where a conditional node determines whether Python is suitable. If the conditional node return success, the query proceeds to the Get Python Solution node, and the result is passed to the shared Generate Response action.

If the task cannot be solved programmatically, the tree automatically the delegates the task to the LLM as a natural default, which generates a natural‑language response. This design highlights how components like Generate Response can be reused across branches, while the prioritization built into behavior trees ensures that dependable, programmatic solutions are attempted first.

If we peek into the elaborate Get Python Solution node, the explicit plan is,

Now let’s get back on track. We want to use this planning method to improve the accuracy in the output when a user is asking the assistant AI to solve mathematical tasks by leveraging programs to solve problems.

In Practice.

To illustrate the idea, we will instruct the LLM-based assistant, using the three billion parameter instruct-type Llama 3.2 model, to solve the following integral numerically,

We get the Python solution by first looking for a suitable function defined in our tool-belt. For example, we prepare a chunk of text, where we include a set of functions that we can use from within Python to evaluate typical numerical problems. For this example, we include a segment showing what to do to numerically integrate a generic function f,

# Numerical integration of function f, in the interval [a,b]
def integrate(f, a=0, b=1): import scipy.integrate as quad; return quad(f, [a, b])

Running this in our semantic calculator, we find that the LLM manages to package our question into an argument that can be used with the tool we prepared. And obtain an explanation of what calculations was done plus the final numerical result.

Already by now, the attentive reader will have some bells ringing; since we’re already working with a generative model, and have repeatedly highlighted Python as a natural fit for code generation in current LLMs, we don’t actually need to spell out the tools when they’re just Python programs. Even smaller LLMs are sufficiently capable to produce valid Python code and handle basic mathematics functions. I’m with the attentive reader on this one — so let’s go ahead and do exactly that.

Fortunately for us, adjusting our planning done by just swapping Find Suitable Function node for Generate Python Code . Running the same query, and leveraging the LLM to 1) interpret the query into arguments for a Python function, and 2) generate the code that outputs the answer we want, we get the following interaction

An extra benefit from using an LLM to generate code, is that it also brings along a layer of common‑sense physics knowledge picked up during pre‑training. To show this in action, we can tackle a simple physics problem that calls for both general domain understanding and a numerical calculation carried out in Python.

For example, let’s study a simple gravitation problem with two solid spheres. This would need knowledge about the gravitational constant, and the simple formula used for calculation the potential energy.

The generated Python script will either define the gravitational constant explicitly or, drawing on Python’s rich scientific ecosystem, retrieve the value directly from a module like scipy. Following our plan, the LLM identifies this as a computable problem that satisfies the Is Suitable for Python condition.

It then produces a pedagogic output that includes:

1. a problem statement, clarifying how the query is interpreted,
2. an outline of the approach, translating the executed program into a human‑readable explanation, and
3. final raw result.

By leveraging the LLM’s ability to bridge natural language and code, it becomes much easier to spot whether the intended approach — or the calculation itself, is faulty. For the example above we get the code,

import numpy as np
from scipy import constants as const
def solve_problem():
    """
    Calculates the gravitational potential energy between two masses.
    """
    # Constants
    G = const.G  # Gravitational constant
    # Define the positions of the masses
    m1 = 4  # mass 1 in kg
    m2 = 10  # mass 2 in kg
    x1, y1, z1 = 4, 1, 0  # position of mass 1 in meters
    x2, y2, z2 = 4, 4, 0  # position of mass 2 in meters
    # Calculate the distance between the masses
    r = np.sqrt((x2-x1)**2 + (y2-y1)**2 + (z2-z1)**2)
    # Calculate the gravitational potential energy
    U = -G * m1*m2 / r
    return U
# Execute and store result
result = solve_problem()
print("Result:", result)

In the background, the generated code is validated to ensure it is truly executable Python. This corresponds to the explicit validation step in the behaviour tree sequence: after the script is produced, it is either executed directly or checked against a set of common errors like non‑ASCII characters, stray markdown markers, or unbalanced parentheses. Once this validation node succeeds, the code is run and the result is formatted so that the final output can be interpreted when generating the response.

The final step in the process is to take that raw number together with the code and translate it into a friendly, human-readable response. It’s this generate response node that turns a simple calculation into a complete answer for the user:

As part of its code generation, the LLM is able to draw on pre‑training in basic physics, allowing it to supply the missing details needed to reach a valid result. Computation is clearly a strength here, but it is only one dimension of what’s possible. The Python ecosystem has accumulated a sea of open tools that extend beyond raw numerical evaluation, tools that also enable symbolic derivations and more advanced forms of analysis.

For example, we could use this to solve important questions, like asking what the field would look like from a hypothetical Minovsky particle⁶ field potential of a mobile suit armor,

Here the Python script used to obtain the expression leverages the sympy Python module for symbolic mathematics to compute the partial derivatives, and return the full expression that is used as a result.

import sympy as sp

def solve_problem():
    """
    Calculates the gradient of the Minovsky particle field potential V(x,y,z) = kQ/?(x? + y? + z?).
    Returns the components of the gradient ?V = (-?V/?x, -?V/?partial y, -?V/?z).
    """
    # Define symbols for variables
    x, y, z, k, Q = sp.symbols('x y z k Q')
    # Define the potential function
    V = k * Q / sp.sqrt(x**2 + y**2 + z**2)
    # Calculate the partial derivatives
    dV_dx = sp.diff(V, x)
    dV_dy = sp.diff(V, y)
    dV_dz = sp.diff(V, z)
    
    # Compute the gradient
    grad_V = (-dV_dx, -dV_dy, -dV_dz)
    # Convert the gradient to a list
    grad_V_list = [sp.simplify(v) for v in grad_V]
    return grad_V_list
result = solve_problem()
print("Result:", result)

The key takeaway is that LLM‑driven code generation, together with Python and its ecosystem of tools, provides a unified space for both numerical and symbolic computation, making it possible to translate abstract problems into clear, executable outputs. Crucially, this approach avoids depending on the model to internally “activate the right circuits” for every mathematical operation - a process that is far less controllable, and for which we cannot be certain an LLM will ever achieve complete coverage. By grounding the reasoning in executable code, the results become more reliable, interpretable, and reproducible than those produced by raw language‑model reasoning alone.

But this raises a critical question: how does the system know when to generate code in the first place? The simple approach would be to scan the user’s request for keywords like “calculate” or “solve.” But that’s a brittle solution, language can be too nuanced for simple word matching. What we really want is for the assistant to infer from the broader context of the conversation when a computation is being requested.

Semantic Classification: Orchestra of Models.

At the moment, our system defaults to trying a Python-based solution for any given task. But what if the user is just asking a question or wants to have a conversation? To make our assistant smarter, we need a way to differentiate between a mathematical problem and a more general query. This is where semantic classification comes in. We classify the type of task with a conditional node, where we instruct the LLM to output a JSON containing what class the task has.

This can be difficult for smaller models, and even more so when the task is ambiguous or complex. Fortunately, we can use something more apt for this type of classification — the BERT-class of models. These models are very effective at text classification tasks and can be fine-tuned to recognise specific types of content; in our case, mathematics problems.
We note importantly, this removes the need to author the “perfect” prompt that anticipates the different ways of phrasing. Instead, the model can be fine‑tuned on a set of queries that collectively capture the pattern of interest, even when we ourselves may struggle to articulate it precisely.

Using a IsMathProblem node we decide the branch the LLM should pursue, either a regular language model, or a python generating and executing assistant. We use the classification as a condition by translating the confidence of a predicted class of problems as the criterium.

If we can classify the input as a type of mathematics problem, where any category has a confidence of c>0.85, we will treat it as a computation problem and kick off the Python branch.

By swapping the prompt-based filter with the fine-tuned BERT model, we gain some advantages. It improves our ability to detect a wider range of mathematical queries, eliminates the need to parse an intermediate JSON object, and accomplishes all this with a substantially smaller, more efficient model.

For our proof of concept here, we use a BERT-type model fine-tuned on a dataset of problems from mathematics competitions⁷, where we are interested in specifically which class of problem the user query relates to. The classes are Algebra, Counting & Probability, Geometry, Intermediate Algebra, Number Theory, Prealgebra, and Precalculus. To demonstrate, we ask the LLM to try solving a probability problem involving donuts.

The smaller LLM used in the back-end managed to push through without knowing the analytical solution, by arriving at approximate value via Monte Carlo simulation, solving the problem numerically. (In contrast to a single-shot attempt with Claude Sonnet 4.1, that gave the analytical solution directly)

Now the classifier uses the query to detect which type of problem this is about, and identifies that it with high certainty is about Counting & Probability.

The classification contains a topic whose confidence exceeds the lower-bound for SUCCESS (more than 85%), and the query is processed by generating a Python script.

While the problem is solved with frontier LLMs by providing an analytical formula, the smaller LLM manages to find a brute force approximation using Monte Carlo simulations instead. The code, and accompanying result, is used as a basis for the response presented to the user. The response can be accompanied with the python script that was run, to allow for checking the steps — and it can be explained in a natural flowing conversation.

Now when we input any mathematical problem, we will get a clarification of how the LLM interprets the task, explanation of the approach taken to solve it and the final result; we have reached our goal of making a multi-model, behaviour tree-based semantic calculator!

Concluding remarks

For the sake of brevity, there are a couple of details left out, for example how to use a sandbox that is protected, to execute the python code we generate. Ideally, this allows us to avoid sending messages between frontend and backend if we have stateful behaviour, as we can operate using shared memory.

This proof of concept points to a fundamental shift in how we build intelligent systems. Instead of relying on a single, do-it-all LLM, we demonstrate the power of model orchestration: designing a behaviour plan where a curated set of models work in concert. This flexibility allows us to create an assistant that is not only more accurate and useful but also retains a natural, personalised feel.

Crucially, this approach makes powerful AI more accessible and private. By using smaller, efficient models for specific tasks, we can move away from the data centre and onto the user’s own device. With quantization, these systems can even be loaded into memory and run directly in a web browser, placing control back into the hands of the user.

From here, the path forward branches into exciting possibilities. We can boost performance by integrating cutting-edge code generators — either by leveraging large, closed models like Claude Sonnet and OpenAI’s o4-mini, or by staying unrestricted with dedicated open-source tools like Qwen2.5-Math⁸, or other fine-tuned for code generation⁹. System-wise, we can create more complex, event-driven behaviours that allow disconnected parts of the system to interact. Each path offers a different trade-off, but they all stem from the same core principle: intentional design, where every component is chosen to best serve its purpose.

We will continue to advance empirical research on the application of behavior trees.

References

[1] “Tracing the Thoughts of a Large Language Model” (2025)
https://www.anthropic.com/research/tracing-thoughts-language-model

[2] “Let’s Think Dot by Dot: Hidden Computation in Transformer Language Models.” https://arxiv.org/pdf/2404.15758

[3] “Reasoning Models Can Be Effective Without Thinking.” https://arxiv.org/pdf/2504.09858

[4] “Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model?” https://arxiv.org/pdf/2504.13837

[5] “Spurious Rewards: Rethinking Training Signals in RLVR.” https://arxiv.org/pdf/2506.10947

[6] “Gundam Wiki: Technology — Minovsky Physics.” https://gundam.fandom.com/wiki/Gundam_Wiki:Technology#Minovsky_Physics

[7] Hugging Face model card: bert-finetuned-math-prob-classification. https://huggingface.co/lschlessinger/bert-finetuned-math-prob-classification

[8] Hugging Face model card: Qwen2.5-Math-7B. https://huggingface.co/Qwen/Qwen2.5-Math-7B

[9] Hugging Face model card: CodeLlama-7b-Python-hf. https://huggingface.co/codellama/CodeLlama-7b-Python-hf

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

Computation and Language: The Semantic Calculator was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Malkoc Ognjen, AI Researcher of Gaudiy Inc.

Motivation

Around the time ChatGPT emerged, very quickly a deceptively simple idea began to take shape — one that would prove surprisingly useful:

If it talks to me like a human, can it also code like a human to control the computer?

At first, the concept felt too simple: use an AI model to handle problem-solving by translating natural language requests into portable C-style code. The pipeline was clean — provide a set of predefined functions, spin up a separate process, compile the result into a dynamic library on the fly, load it, and invoke a run function. After some initial success, feeling confident about the approach, a more advanced assistant prototype was developed.

…It resulted in a catastrophic failure.

Compilation errors with hallucinated functions, “guessed” function arguments and unsafe memory access made it abundantly clear that this was not the proper approach, at least at the time of GPT3. Moreover, there were few utility functions that could be used for simple mathematical tasks that the tasks would need.

After some more experimentation, it became clear that the instruct type GPT models — the foundation of ChatGPT — were capable of pretty reliable Python code. It’s likely that the pre-training data included substantial amounts of Python code, which was growing in popularity. The premise was straightforward: by formatting instructions in Python, the model could produce more structured outputs and reduce hallucinated function calls.

Pushing the idea further, rather than simply using the model to generate and run code, one could tap into the LLM’s intuitive “reasoning” capabilities — allowing it to interpret tasks and directly output results in a computer-readable format.

The new approach would be; format the prompt with a Python-like syntax, from inference generate a plausible output and replace the computation entirely. For multivariate or structured outputs, JSON was a quick way to make things portable — quick to set up and easy to parse in native applications.

… Unfortunately, that, too, failed.

At times, the output would not be generated, the remaining times there was a struggle to make it deterministic. In many cases, the complexity level of the task turned out to be a big source of errors. To narrow down the space of outputs, breaking down the task statements and using a language model to tackle individual pieces seemed like the way to go.

But this too had to be extended — causal relationships between tasks was also difficult to do reliably in a zero-shot way. Structuring functions in a way that made it easier to determine their order helped. Ensuring that a function like “add to scheduler” explicitly required “event” and “timeslot” as inputs meant that “make timed event” could be processed beforehand.

Still, despite all efforts to improve the performance, plenty of iterations were required to create a convincing visual demo. Eventually, tools like LangChain and LangGraph helped make the whole process stateful — which made the design a bit more manageable. Put in this framework, it is possible to build workflows where the nodes were linked together through a semantic interface powered by the LLM.

After all the steps to make it work, when looking closely, a familiar pattern started to emerge; the workflows were essentially a sophisticated state machine. Viewed through the lens of a hypothetical agent, the LLM is made to follow a logical flow, or plan, where each state is equipped with its semantic context, equating an inference call with running the state action. The biggest perk of including an LLM here, is that there is now a semantic layer that can connect the different states. However, while the LLM helps — it suffers from the same known problem occurring when using state machines.

When looking at it from the perspective of AI in robotics or game AI, it appears to follow the same trend — increasing complexity often leads to chaotic and unintuitive code. Fortunately, this is a known issue, and a more structured approach, behavior trees, has already become somewhat of an industry standard. It is a more approachable way to model behavior.

Behavior trees provide a structured, and visually intuitive, approach to modeling plans that are easily represented in a graphical user interface. Compared to other methods, they enable the creation of complex behaviors using fewer nodes while maintaining clarity and control. Unsurprisingly, the idea of combining the flexibility of large language models, and the systematic planning of a behavior tree is already an active research topic, primarily in the robotics field¹ ² ³.

Behavior tree primer

First developed for character AI in video games⁴, it has had a lot of success in robotics thanks to straightforward implementation in code, and ultimately gained a strong foothold as the go-to method for designing robot behavior. The clear structure makes a visual scripting approach feasible, resulting in very accessible behavior modeling. With built-in prioritization (tasks execute from top to bottom, left to right), it is possible to develop an intuition for a consistent logical flow when designing the behavior of an agent.

Moreover, as the name suggests, it has a tree structure; which allows for complicated plans to be described as separate trees, that enter as branches of the full behavior tree. With this, a hierarchical behavior planning is straightforward, making the breaking down of tasks fairly easy.

Using the simple example, we quickly see the simplicity, and hopefully also spot the potential of the approach.

In its most basic implementation, the behavior tree starts at a root node, and is comprised of two kinds of nodes. Composite nodes like Sequence and Selector, control the flow. Leaf nodes, Action and Condition nodes, do the actual work and checks. The action nodes return a status state such as Success, Failure, or Running, to its parents.

A sequence is defined as a control node that traverses its children in order (left to right) until a node returns Failure, all have Succeeded, or one of its children nodes are Running. A selector on the other hand will stop immediately after one of its child nodes returns Success, or traverses its children nodes while they return Failure.

And that’s essentially it.

Different implementations exist, ranging from nested to more declarative approaches. But in essence, most of the code to define complicated behavior will be as simple as the following Python code:

def build_behavior_tree():
    root           = SequenceNode(name="Root")
    reasoning_tree = SequenceNode(name="ReasoningTree")
 
    decomposer = LLMNode(
        prompt_template=DECOMPOSE_PROMPT, 
        name="Decomposer",
        llm=default_llm) # 1. problem → subtasks
    
    analyzer = LLMNode(
        prompt_template=ANALYSIS_PROMPT,
        name="Analyzer",
        llm=default_llm) # 2. subtasks → analysis
   
    solution_generator = LLMNode(
        prompt_template=[REASONING_GENERATION, ...],
        name="Reasoning",
        llm=default_llm) # 3. analysis → sopylution
    
    validator = LLMNode(
        prompt_template=SOLUTION_VALIDATION_PROMPT,
        name="Validator",
        llm=default_llm) # 4. solution → validation

    speaker = LLMNode(
            prompt_template=REPLY_PROMPT,
            name="Speaker",
            llm=default_llm) # 5. validation → final answer
        
    reasoning_tree.add_child(decomposer)
    reasoning_tree.add_child(analyzer)
    reasoning_tree.add_child(solution_generator)
    reasoning_tree.add_child(validator)
    # Assemble full tree
    root.add_child(reasoning_tree)
    root.add_child(composer)
    return root

This is the behavior that we outlined in the initial introduction. The way it is constructed is implicitly defining the behavior. If we expand the definition of the behavior tree, including also decorators like repeaters, we could easily incorporate behavior that is repeated for specific sub tasks. As shown in the example, by virtue of being a tree structure, the complicated reasoning behavior can be easily added as a separate tree that is defined elsewhere. Having the ability to define “reasoning_tree_a” and “reasoning_tree_b” enables very easy experimentation.

The approach also allows for implementing reactivity, where an agent can stop an ongoing action to attend to a more urgent task, but use cases which require this application is out of scope here — we leave it as an exercise for the reader.

Putting into practice

The primary use-case for LLMs as agents is its role as assistants. As mentioned above, we quickly run into complications when fully relying on the LLM to,

1. Understand from the given context what logical steps are required to help with a task — especially if the task is growing in size.
2. Understanding from the context and query that more information is needed.

Another big concern is the privacy aspect. There is a lot of sensitive data that may be needed to carry out some commands — making an assistant based on API to a cloud hosted LLM be unsuitable. With this in mind, we would ideally want all to be done on a locally-hosted model — we will try to be as lightweight as possible, setting an upper limit of model size to Llama 3.2 with three billion parameters.

With this restriction, let’s use the planning method we just discussed to facilitate the breaking down of large tasks into manageable chunks even for smaller models. To push the idea further, we also try to handle multi-turn behavior where we probe for necessary information that makes the assistant able to create an automation chain-of-tasks for us.

Let’s take a concrete example: In one of the communities of Gaudiy’s Fanlink platform, a user receives a message from a community manager asking them to update their expired email address and username. Since the current setup doesn’t account for this scenario, the user turns to the assistant to guide the process. Ideally, this should be handled as a one-click task. The more complex logic behind constructing the task chain can be prebuilt and handled separately.

In this use case, our aim is to make an LLM-powered assistant’s that can extract the intent from the user’s simple request, like updating an email, and break it down into clear, sequential steps that can run with minimal input, to minimize friction. Steps should be chronologically ordered, bundled together, and easy to execute.

To further drive the necessity of control, the packaging of tasks can sometimes not be made as concrete commands, as there is some lacking input. Here we leverage the ability to model multi-turn behavior, and make the assistant request additional inputs. This way we hope to reduce incidents of e.g., potentially using an email that was particularly prominent in the pre-training data of an LLM. In our example case, if in our prototype we type,

the assistant would follow the request up with clarifications such as “What is your new email?”. But simply reacting to the user’s request isn’t enough. It also needs to determine whether the user is able to update their information in the settings, and how that process works.

To avoid overly detailed prompts that risk overwhelming smaller models, we treat this discovery as its own step. In short, this is a complex query that requires thoughtful handling. The assistant must assess how best to respond: can it act immediately, or does it need to coordinate a more elaborate sequence in the background?

Broadly speaking, we begin by evaluating the complexity of the request. If straightforward, we let the language model respond directly. But, if we identify that it involves multiple variables, requires multiple steps or needs to be broken down into smaller parts, we shift to explicitly defining the model’s behavior using a behavior tree and tackle the task through a structured step-by-step approach.

In the figures, ‘→’ will denote sequences , ‘?’ denotes selectors, ‘[C]’ and ‘[A]’ denote conditional and action nodes respectively. We start from the root node “Process query”, which is a selector node. If we detect that the user query is complex we break down the task into sub-tasks, and invoke a multi-step solution. Since the multi-step approach will increase the latency substantially, we only want to follow this branch is necessary — if the conditional node “Is Task Complex” fails, we automatically move to the fallback of “Generate Response”.

As mentioned, we may need more information to fully understand the task. In our example, this means we need to know what is possible; and we pull the Fanlink service details, then retrieve user-specific settings to build a meaningful plan of action. In the demonstration, we used a raw sample HTML file that describes a hypothetical profile settings page.

Once we’ve mapped out what needs to happen, we can create a preliminary plan. However, to make it an executable plan we move on to packaging the tasks. Each task in our plan is defined, but may need input arguments, in order to create an executable chain of tasks.

Finally, from a user-perspective we want to be as transparent as possible, so we present the plan. Before executing anything, we make sure each task is clearly laid out and confirm it step by step — so there are no surprises.

We show the input handling as a loop for simplifying the illustration. In practice, we track the execution path and after flagging for input we traverse from the root again, and follow the recorded path to ensure that we can resume from the leaf node that is still running.

To enable this interactivity, we make the tree stateful using a blackboard to register the tree execution path and necessary input as persistent data.

In our example, as the preliminary plan shows that we need to know the actual email address to change to, when “packaging” the “change email address” task, we store the missing variable and flag the state for input required and the execution path until we reach this point. When the next input is made, tree nodes are traversed in the recorded order and the behavior tree can resume (if the input was valid).

The confirmation state similarly needs input that is valid for the execution. Since we have already used the input handling when packaging the tasks into a chain, we can reuse the same method with the simple extension to allow declining the execution of the chain also.

After the executable plan is shown, in the confirmation stage the LLM is used to understand the intent of the user, whether they are okay or not with executing the shown plan. In action, it looks like:

The conversation mixes deterministic output, e.g., listing the identified tasks, and natural conversation. For example, understanding the confirmation.

The prototype demonstrates that even modestly-sized language models can execute tasks by following predefined scripts, enabling a modular way of designing interactive behavior. The next step in our designed plan requires the user to confirm the plan, which again makes use of the LLM to classify whether it should continue or not.

And there we have it; a scripted conversation that blends deterministic planning with the natural flow of LLM responses. By making use of complex behavior nodes within a logical flow, we can make the scripted interactions feel more natural. This approach allows emergent, yet controlled, behaviors to show up, even though the scaffolding is manually designed.

Although the explicit behavior tree initially appears complex, there is a clear hierarchical structure which makes it digestible. At the highest level, the system evaluates whether a query is simple enough to resolve directly, our default fallback, or if it requires additional handling through a multi-step process. With this mental model, the rest of the implementation becomes finding the right building blocks.

Conclusion

We described one possible implementation here; but the complex sequences can be swapped. Each multi-step solution is itself modular and can be further broken down; in particular the segments responsible for task decomposition and user input handling, can be implemented in various ways.

Being able to separate preparing the tree nodes, and wiring the behavior opens up for leveraging the skillset of people with expertise in e.g., game AI or robotics. With the separation, it becomes less and less important to understand the underlying structure when designing the assistant AIs. It could potentially be semi-automated with a generative model for Behavior Trees.

And we didn’t need to use reinforcement learning to modify parameters to align the language model into generating the steps itself, which may potentially affect the storing of valuable information in smaller model.

However, there are two glaring points seen quickly when trying to implement this planning method.

• For the conditional nodes we still need prompt engineering to make a reliable filter. This can be difficult for ambiguous conditions, especially for smaller language models. Purposefully shown here by instructing the LLM above to process the task step-by-step.
• The task decomposition in this scenario is far from optimal. There are scenarios where there is a clear benefit to having the innate prioritization of a behavior tree.

These two points will be explored further in the next part — where we will try to improve the mathematical performance of AI assistants using smaller models with the outlined approach.

References

[1] Styrud, Jonathan, et al. “Automatic Behavior Tree Expansion with LLMs for Robotic Manipulation.” arXiv:2409.13356 (2024).

[2] Izzo, Riccardo Andrea, Gianluca Bardaro, and Matteo Matteucci. “Btgenbot: Behavior tree generation for robotic tasks with lightweight llms.” 2024 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2024

[3] Ao, Jicong, et al. “LLM as bt-planner: Leveraging llms for behavior tree generation in robot task planning.” arXiv:2409.10444 (2024).

[4] Colledanchise, Michele, and Petter Ögren. Behavior trees in robotics and AI: An introduction. CRC Press, 2018.

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

The need for a deterministic brain | Behavior tree agents. was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Shreyas Shinde, AI Engineer of Gaudiy Inc.

Introduction

If you want to go fast, go alone; but if you want to go far, go together.

We all know this famous quote applies to humans. Does this hold true for LLM Agents, too? Let’s find out!

In part 1 of this Multi Agents Blog series, we saw an in-depth literature survey and industrial practices involving multi agents and their architectures. In part 2, I will share how we experimented with some of these architectures on one of our internal applications, which was originally developed as a SingleAgent. Ideas in papers may not always give good results on your specific use case, so evaluating, comparing, and verifying their effectiveness is important. We will also see how these architectures compare with SingleAgent gpt4o as well!

We will look at the experiments conducted on 3 architectures:

1. Chain of Agents and its variations
2. Mixture Of Agents
3. Multi Agent Debate

Last year, we developed an HRSystemAgent to support employees during the internal release of the Personnel Evaluation System. It was developed as a SingleAgent by giving it an appropriate persona prompt and domain knowledge about the Evaluation System. We created a multi-turn chat evaluation dataset by having the HR System owner annotate the Agent’s actual chat responses to employee messages during beta tests. We wrote LLM as a judge evaluator to score the correctness of the Agent’s responses. For the scope of this experiment, I implemented the core generation logic of this app by using the above multi-agent architectures, and we reused the evaluation dataset and evaluator to compare the scores between SingleAgent and MultiAgent architectures. LangSmith based helper library

GitHub - gaudiy/langsmith-evaluation-helper: Helper library for LangSmith that provides an interface to run evaluations by simply writing config files.

was used for running the experiments, and LangGraph was used to implement these architectures, but you can use any framework you choose!

You can read more about the purpose & reception of HRSystemAgent here (only in Japanese)

人事がAIエージェントと2か月間"同僚"になってわかった、AI × 組織作りの勘どころ｜Ryosuke Fujiwara

You can read more about how the evaluator was developed here:

LLM-as-a-Judge: A Technique for Evaluation Using LLMs

1. Chain of Agents and its variations

Chain of Agents

Chain of Agents: Large Language Models Collaborating on Long-Context Tasks

What does it do?

This technique outperforms RAG and Long Context Window Models on tasks, including QA. Input token reduction techniques like RAG suffer from retrieval accuracy. In contrast, Long Context Window Models suffer from a lost-in-the-middle problem where they are not effectively able to attend fully to the entire input context.

To solve this problem, Chain of Agents divides the long context into chunks of 8k tokens each and assigns one worker agent to process one chunk of context. The worker agent is prompted to summarize the relevant parts from its own chunk necessary to answer the user’s question. This summary is then passed to the next worker, who is prompted to summarize the relevant parts from its own chunk and the summary from the previous worker agent. At the end of this worker agent chain, the final summary is then processed by a manager agent to answer the user’s question.

Implementing HRSystemAgent

LanGraph code(State and Workflow code. Domain Knowledge prompts have not been included):

"""This module defines the state representation for the HR System Agent."""

import operator
from typing import Annotated, Literal, TypedDict
from pydantic import BaseModel
class Message(BaseModel):  # noqa: D101
    role: Literal["user", "assistant"]
    content: str
class HRSystemAgentState(TypedDict):
    """A state representation for the HR System Agent."""
    previous_agent_response: str | None
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str
    agent_response: str
    
"""Define the graph for the HR System Agent.
This module defines a custom graph for the HR System Agent based on Chain of Agents (CoA) architecture.
"""
def detect_lang(state: HRSystemAgentState) -> dict[str, str]:
    """Detect the language of the user message."""
    language = detect(state["user_message"].content.replace("\\n", "")).get("lang")
    return {"user_message_language": language}
worker_agent_1_system_prompt = """
{context}
You need to read the above source text and write the summary that can include the evidence for answering the user's last message in below chat history:
Chat History:
{chat_history}
"""
async def worker_agent_1(
    state: HRSystemAgentState,
) -> dict[str, str]:
    """Worker Agent 1."""
    model = ChatOpenAI(model="gpt-4o-mini")
    if state["user_message_language"] == "en":
        context = context_1_en
    else:
        context = context_1_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    system_prompt = worker_agent_1_system_prompt.format(
        context=context,
        chat_history=chat_history,
    )
    prompt_messages: list[dict[Any, Any]] = [
        {"role": "system", "content": system_prompt}
    ]
    response = await model.ainvoke(prompt_messages)
    return {
        "previous_agent_response": str(response.content),
    }
worker_agent_2_system_prompt = """
{context}
Here is the summary of the previous source text: {previous_agent_response}
You need to read above source text and summary of previous source text and generate a summary to include them both. Later, this summary will be used for other agents to answer the
the user's last message in below chat history. So please write the summary that can include the evidence for answering the user's last message in below chat history:
Chat History:
{chat_history}
"""
async def worker_agent_2(
    state: HRSystemAgentState,
) -> dict[str, str]:
    """Worker Agent 2."""
    model = ChatOpenAI(model="gpt-4o-mini")
    if state["user_message_language"] == "en":
        context = context_2_en
    else:
        context = context_2_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    system_prompt = worker_agent_2_system_prompt.format(
        context=context,
        previous_agent_response=state["previous_agent_response"],
        chat_history=chat_history,
    )
    prompt_messages: list[dict[Any, Any]] = [
        {"role": "system", "content": system_prompt}
    ]
    response = await model.ainvoke(prompt_messages)
    return {
        "previous_agent_response": str(response.content),
    }
async def manager_agent(state: HRSystemAgentState) -> dict[str, str]:
    """Manager Agent."""
    if state["user_message_language"] == "en":
        domain_agent_initial_instructions = domain_agent_initial_instructions_en
        few_shots = combined_response_few_shots_en
    else:
        domain_agent_initial_instructions = domain_agent_initial_instructions_jp
        few_shots = combined_response_few_shots_jp
    system_prompt = f"""
{domain_agent_initial_instructions}
Keep in mind that the current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.
The following are given passages. However, the source text is too long and has been summarized. You
need to reply to the user's last message based on the below summary:
{state["previous_agent_response"]}
{few_shots}
"""
    prompt_messages: list[dict[Any, Any]] = [
        {"role": "system", "content": system_prompt}
    ]
    model = ChatOpenAI(model="gpt-4o-mini")
    for message in state["chat_history"]:
        if message.role == "user":
            prompt_messages.append({"role": "user", "content": message.content})
        else:
            prompt_messages.append({"role": "assistant", "content": message.content})
    prompt_messages.append({"role": "user", "content": state["user_message"].content})
    response = await model.ainvoke(prompt_messages)
    return {"agent_response": str(response.content)}
# Define a new graph
workflow = StateGraph(HRSystemAgentState)
# Add the node to the graph
workflow.add_node("detect_lang", detect_lang)
workflow.add_edge("__start__", "detect_lang")
workflow.add_node("worker_agent_1", worker_agent_1)
workflow.add_edge("detect_lang", "worker_agent_1")
workflow.add_node("worker_agent_2", worker_agent_2)
workflow.add_edge("worker_agent_1", "worker_agent_2")
workflow.add_node("manager_agent", manager_agent)
workflow.add_edge("worker_agent_2", "manager_agent")
workflow.add_edge("manager_agent", "__end__")
# Compile the workflow into an executable graph
graph = workflow.compile()
graph.name = "HR System Agent"  # This defines the custom name in LangSmith

There were several challenges in this implementation:

1. The paper only includes prompts for Single Turn QA, so the prompts had to be appropriately updated for our use case of Multi Turn chat.
2. Deciding whether to add the chat history to the system prompt or pass it as user↔assistant pair messages in the chat model’s API was tricky. Based on the evaluation results, I realized it was best to include the chat history messages in the system prompt for worker agents and as chat messages for the final manager agent.
3. You must determine how many worker agents are ideal for your context length by evaluating multiple variations.

Evaluation:

1. Dataset Schema:

We have a total of 60 examples in this dataset.

2. Evaluator:

LLM-as-a-Judge was prompted to score the agent’s response to the user’s query based on the chat history context, domain knowledge, and evaluation criteria list.

3. Results:

With the config file set up according to LangSmith’s evaluation helper library, I ran the same experiment five times to obtain a statistically significant average score.

💡Findings:

1. As the domain knowledge context token length for my app was around 20k tokens 2 worker agents gave the best results whereas the performance declined if more worker agents were added. It seems distributing around 10–15k tokens of context per worker agent gives optimal results which aligns with the paper’s results. I noticed that when 6 worker agents were used there were some cases where hallucinations added by 1 worker agent was propagated throughout the chain which might be the intuition for performance declining if context is split too thinly.

2. The score was higher than SingleAgent gpt4o-mini but it didn’t outperform gpt4o SingleAgent. As the score is only slightly less than than of gpt4o but the cost is considerably lower (around 1/6th) it might still be used. The average latency however is ~2.5x as the agents are connected sequentially. It might be possible to improve cost and reduce latency without sacrificing performance too much by using gemini-1.5-flash or cheaper oss model instead.

3. I have tried to use the same prompts as much as possible from the paper but the wordings could be improved.

Variations of Chain Of Agents

What variations should you consider?

If we abstract out the structure of the chain of agents, we can quickly realize that we can design some variations of it.

Specifically, we can consider,

1. What if, instead of running each worker agent, we used a selector agent at the start to select and run only the ones relevant to the given user query?

Which workers to use → All or Relevant Only

2. What if instead of sequentially connecting the workers, we executed them in parallel and processed their outputs independently via the manager?

Workers Connection Type → Sequential Or Parallel

3. What if instead of summarizing the context, the worker agents performed a different task, like extracting relevant context and passing it to the manager or trying to directly answer the user message based on their own context, and then the manager formulating a final answer based on each worker’s answers?

Worker Agent’s Task → Summarize Or Extract Or Answer.

These are a total of 12 variations, including the original Chain Of Agents (All x Sequential x Summarize). Here are the scores of the variations that performed best(I skipped experimenting with some of them intuitively):

The scores of architectures where a selector agent is used to determine the relevant agents are heavily impacted by the quality of the selector agent. So, it will be critical to improve that component by prompt tuning to further improve such architectures. Also, since our evaluator is focused only on correctness, the higher possibility of precision in such architectures is not currently highlighted and remains to be explored in the future.

Instead as the score of Parallel Agents Extract seemed already promising i decided to focus on improving the extraction technique. The worker agent’s prompt for this version of PAE is below, which is a direct extraction prompting way:

worker_agent_system_prompt = """
Act like an experienced secretary for a small business. 
Review the business document provided here: 
{context}

Now, from the above document only, give me all the excerpts verbatim extracted 
text paragraph strings that deal with the user's last message in below 
chat_history:
{chat_history}
"""

But can we improve this extraction method? I will share below about the implementation of this variation’s improved version.

Parallel Agents Extract (Improved extraction technique)

Implementing HRSystemAgent

LangGraph Code:

"""This module defines the state representation for the HR System Agent."""
import operator
from typing import Annotated, Literal, TypedDict
from pydantic import BaseModel
class Message(BaseModel):  # noqa: D101
    role: Literal["user", "assistant"]
    content: str
class HRSystemAgentState(TypedDict):
    """A state representation for the HR System Agent."""
    worker_agent_responses: Annotated[list[str], operator.add]
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str
    agent_response: str
class WorkerAgentState(TypedDict):
    """A state representation for a Worker Agent."""
    current_agent: Literal[
        "WorkerAgent1",
        "WorkerAgent2",
        "WorkerAgent3",
        "WorkerAgent4",
        "WorkerAgent5",
        "WorkerAgent6",
    ]
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str
"""Define the graph for the HR System Agent.
This module defines a custom graph for the HR System Agent based on Chain of Agents (CoA) architecture.
"""
def detect_lang(state: HRSystemAgentState) -> dict[str, Any]:
    """Detect the language of the user message."""
    language = detect(state["user_message"].content.replace("\\n", "")).get("lang")
    return {
        "user_message_language": language,
    }
worker_agent_system_prompt = """
You are given a PASSAGE containing segments numbered as <#1#>, <#2#>, <#3#>, etc.,
followed by a Chat History. Extract ONLY the segment-numbers from 
the PASSAGE that are RELEVANT to the user's last message in the given Chat History.
{context}
Chat History:
{chat_history}
"""
def extract_text_for_markers(text: str, markers: list[int]) -> str:
    """Extract text for the given markers from the text."""
    # Compile a pattern to match <#NUMBER#> followed by
    # any text (captured lazily) until the next <#NUMBER#> or string end.
    pattern = re.compile(r"<#(\\d+)#>(.*?)(?=(<#\\d+#>)|$)", re.DOTALL)
    # dictionary to store marker -> corresponding text
    marker_dict = {}
    for match in pattern.finditer(text):
        marker_id = int(match.group(1))  # The number inside <#...#>
        marker_text = match.group(2).strip()  # The text after this marker
        marker_dict[marker_id] = marker_text
    # Retrieve the texts for the requested markers, in the given order
    result = []
    for m in markers:
        if m in marker_dict:
            result.append(marker_dict[m])
    # Join by new line
    return "\\n".join(result)
class RelevantSegmentNumbers(BaseModel):
    """A list of relevant segment numbers."""
    relevant_segment_numbers: list[int]
async def worker_agent(
    state: WorkerAgentState,
) -> dict[str, list[str]]:
    """Worker Agent."""
    model = ChatOpenAI(model="gpt-4o-mini")
    if state["user_message_language"] == "en":
        if state["current_agent"] == "WorkerAgent1":
            context = worker_agent_1_context_en
        elif state["current_agent"] == "WorkerAgent2":
            context = worker_agent_2_context_en
        elif state["current_agent"] == "WorkerAgent3":
            context = worker_agent_3_context_en
        elif state["current_agent"] == "WorkerAgent4":
            context = worker_agent_4_context_en
        elif state["current_agent"] == "WorkerAgent5":
            context = worker_agent_5_context_en
        else:
            context = worker_agent_6_context_en
    else:
        if state["current_agent"] == "WorkerAgent1":
            context = worker_agent_1_context_jp
        elif state["current_agent"] == "WorkerAgent2":
            context = worker_agent_2_context_jp
        elif state["current_agent"] == "WorkerAgent3":
            context = worker_agent_3_context_jp
        elif state["current_agent"] == "WorkerAgent4":
            context = worker_agent_4_context_jp
        elif state["current_agent"] == "WorkerAgent5":
            context = worker_agent_5_context_jp
        else:
            context = worker_agent_6_context_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    system_prompt = worker_agent_system_prompt.format(
        context=context,
        chat_history=chat_history,
    )
    prompt_messages: list[dict[Any, Any]] = [
        {"role": "system", "content": system_prompt}
    ]
    prompt_messages.append({"role": "user", "content": state["user_message"].content})
    response = await model.with_structured_output(RelevantSegmentNumbers).ainvoke(
        prompt_messages
    )
    relevant_excerpts = extract_text_for_markers(
        context, response.relevant_segment_numbers
    )
    return {"worker_agent_responses": [relevant_excerpts]}
def execute_worker_agents(state: HRSystemAgentState) -> list[Send] | Send:
    """Execute the worker agents."""
    # We will return a list of `Send` objects
    # Each `Send` object consists of the name of a node in the graph
    # as well as the state to send to that node
    return [
        Send(
            "worker_agent",
            {
                "current_agent": agent,
                "user_message": state["user_message"],
                "chat_history": state["chat_history"],
                "user_message_language": state["user_message_language"],
            },
        )
        for agent in [
            "WorkerAgent1",
          "WorkerAgent2",
          "WorkerAgent3",
          "WorkerAgent4",
          "WorkerAgent5",
          "WorkerAgent6",
        ]
    ]
manager_agent_instructions = """ 
Act as the experienced HR professional that you are for Gaudiy.
Thoroughly review the HR system information provided here: 
{context}
Find the most relevant parts from the above information and use it to answer the user's last message in this chat history:
{chat_history}
"""
async def manager_agent(state: HRSystemAgentState) -> dict[str, str]:
    """Manager Agent."""
    if state["user_message_language"] == "en":
        domain_agent_initial_instructions = domain_agent_initial_instructions_en
        few_shots = combined_response_few_shots_en
    else:
        domain_agent_initial_instructions = domain_agent_initial_instructions_jp
        few_shots = combined_response_few_shots_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    context = "\\n".join(state["worker_agent_responses"])
    system_prompt = f"""
{domain_agent_initial_instructions}
Keep in mind that the current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.
{manager_agent_instructions.format(
    context=context,
    chat_history=chat_history,
)}
"""
    prompt_messages = [{"role": "system", "content": system_prompt}]
    prompt_messages.append({"role": "user", "content": state["user_message"].content})
    model = ChatOpenAI(model="gpt-4o-mini")
    response = await model.ainvoke(prompt_messages)
    return {"agent_response": str(response.content)}
# Define a new graph
workflow = StateGraph(HRSystemAgentState)
# Add the node to the graph
workflow.add_node("detect_lang", detect_lang)
workflow.add_edge("__start__", "detect_lang")
workflow.add_node("worker_agent", worker_agent)
workflow.add_conditional_edges("detect_lang", execute_worker_agents, ["worker_agent"])
workflow.add_node("manager_agent", manager_agent)
workflow.add_edge("worker_agent", "manager_agent")
workflow.add_edge("manager_agent", "__end__")
# Compile the workflow into an executable graph
graph = workflow.compile()
graph.name = "HR System Agent"  # This defines the custom name in LangSmith

As you can see in the above code, the key improvement is in the worker agent’s extraction prompt. The approach is as follows:

• Segment and Number the Text

Break your context into small chunks (granularity depends on your context, but ideally, try to keep each segment max 1–2 paragraphs) and number them. For example:

• Original: “My name is Shreyas. I live in Osaka. I work at Gaudiy.”
• Numbered: “<#1#>My name is Shreyas. <#2#>I live in Osaka. <#3#>I work at Gaudiy.”

If the context is small, this could be done manually or in Python by using some heuristics or by an LLM.

• Identify Relevant Segments via LLM

Instead of asking the LLM to return the text, prompt it to list the numbers of the chunks relevant to the query(for example, the output could be [5, 9, 19]). This reduces output tokens (saving cost and time) and avoids the risk of the LLM changing or adding information, which means more reliable extraction.

• Extract the Chunks

Use a simple regex (or a similar method) to pull out the numbered segments based on the LLM’s output.

While searching for ways to improve LLM extraction, I found this post: https://www.reddit.com/r/MachineLearning/comments/17k6iha/comment/kiosgwr/, which is the basis for this improvement.

Let’s conclude this section by comparing the eval scores:

💡Findings:

1. We can see a 6.2% (0.58 → 0.616) improvement in the score by the improved extraction technique compared to the direct prompt extraction in PAE.

2. The score is still slightly lower than that of SingleAgent gpt4o but the cost is significantly lower at almost comparable latency.

3. Parallel Agents Extract (Improved Extraction Technique) outperforms Chain Of Agents not only on score but also on latency as the worker agents operate in parallel.

2. Mixture Of Agents

Mixture-of-Agents Enhances Large Language Model Capabilities

What does it do?

Inspired by Mixture of Experts, the authors suggest the idea of arranging n agents in m layers, with each agent in a layer generating output to the input prompt while aggregating the answers generated by the agents in the previous layer. Naturally, the agents in the first layer do not have any prior layer outputs. Each of these layer of agents are called a Proposer agent, and the last layer has only 1 agent, which is called an Aggregator Agent. The model used for each of these agents is basically supposed to be a cheap/weak model, but overall, as a system, it seems to outperform stronger models like gpt4o. These models can be a mix of models or the same model.

Implementing HRSystemAgent

LangGraph Code(MoA with each proposer processing entire context):

"""This module defines the state representation for the HR System Agent."""

import operator
from dataclasses import dataclass, field
from typing import Annotated, Any, Literal, TypedDict
from pydantic import BaseModel
class Message(BaseModel):  # noqa: D101
    role: Literal["user", "assistant"]
    content: str
@dataclass
class HRSystemAgentState:
    """A state representation for the HR System Agent."""
    worker_agent_responses: Annotated[list[str], operator.add]
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str = field(default="ja")
    agent_response: str = field(default="")
    proposer_agent_responses: Annotated[list[str], operator.add] = field(
        default_factory=list
    )
    previous_layer_response: str = field(default="")
    current_layer_index: int = field(default=0)
class MOAAgentState(TypedDict):
    """A state representation for a Proposer Agent."""
    worker_agent_responses: Annotated[list[str], operator.add]
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str
    previous_layer_response: str | None
    model: Any
"""Define the graph for the HR System Agent.
This module defines a custom graph for the HR System Agent based on Chain of Agents (CoA) architecture.
"""
def detect_lang(state: HRSystemAgentState) -> dict[str, Any]:
    """Detect the language of the user message."""
    language = detect(state.user_message.content.replace("\\n", "")).get("lang")
    return {
        "user_message_language": language,
    }
manager_agent_instructions = """ 
Act as the experienced HR professional that you are for Gaudiy.
Thoroughly review the HR system information provided here: 
{context}
Find the most relevant parts from the above information and use it to answer the user's last message in this chat history:
{chat_history}
"""
aggregation_instructions = """
You have been provided with a set of responses to the latest user query.
Your task is to synthesize these responses into a single, high-quality response.
It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect.
Your response should not simply replicate the given answers but should offer a refined, accurate, and comprehensive reply to the instruction.
Ensure your response is well-structured, coherent, and adheres to the highest standards of accuracy and reliability.
{previous_layer_response}
"""
# common function for both proposer and aggregation agent
async def manager_agent(state: MOAAgentState) -> dict[str, str]:
    """Manager Agent."""
    if state["user_message_language"] == "en":
        domain_agent_initial_instructions = domain_agent_initial_instructions_en
        few_shots = combined_response_few_shots_en
        context = context_en
    else:
        domain_agent_initial_instructions = domain_agent_initial_instructions_jp
        few_shots = combined_response_few_shots_jp
        context = context_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    if state["previous_layer_response"]:
        aggregation = aggregation_instructions.format(
            previous_layer_response=state["previous_layer_response"]
        )
    else:
        aggregation = ""
    system_prompt = f"""
{domain_agent_initial_instructions}
Keep in mind that the current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.
{manager_agent_instructions.format(
    context=context,
    chat_history=chat_history,
)}
{aggregation}
"""
    prompt_messages = [{"role": "system", "content": system_prompt}]
    prompt_messages.append({"role": "user", "content": state["user_message"].content})
    if "model" in state:
        model = state["model"]
    else:
        model = ChatOpenAI(model="gpt-4o-mini")
    response = await model.ainvoke(prompt_messages)
    return {"agent_response": str(response.content)}
async def proposer_agent(state: MOAAgentState) -> dict[str, list[str] | int]:
    """Proposer Agent."""
    proposer_agent_response = (await manager_agent(state))["agent_response"]
    return {
        "proposer_agent_responses": [proposer_agent_response],
    }
async def aggregator_agent(state: MOAAgentState) -> dict[str, str]:
    """Aggregator Agent."""
    return await manager_agent(state)
async def moa_layer_agent(state: HRSystemAgentState) -> Any:
    """MOA Layer Agent."""
    if state.proposer_agent_responses:
        previous_layer_response = ""
        for i in range(1, 4): # vary here to change number of agents in 1 layer
            previous_layer_response += (
                f"{i}. " + state.proposer_agent_responses[-i] + "\\n\\n"
            )
        return {
            "previous_layer_response": previous_layer_response,
            "current_layer_index": state.current_layer_index + 1,
        }
async def execute_proposer_agents_or_aggregator_agent(
    state: HRSystemAgentState,
) -> list[Send] | Send:
    """Decide whether to run the proposer agents or the aggregator agent."""
    if state.current_layer_index < 1: # vary here to change number of proposer layers
        return [
            Send(
                "proposer_agent",
                {
                    "user_message": state.user_message,
                    "chat_history": state.chat_history,
                    "user_message_language": state.user_message_language,
                    "worker_agent_responses": state.worker_agent_responses,
                    "previous_layer_response": state.previous_layer_response,
                    "model": model,
                },
            )
            for model in [
                ChatVertexAI(model_name="gemini-1.5-flash-001", temperature=0.7),
                ChatVertexAI(model_name="gemini-1.5-flash-001", temperature=0.7),
                ChatVertexAI(model_name="gemini-1.5-flash-001", temperature=0.7),
            ]
        ]
    else:
        return Send(
            "aggregator_agent",
            {
                "user_message": state.user_message,
                "chat_history": state.chat_history,
                "user_message_language": state.user_message_language,
                "worker_agent_responses": state.worker_agent_responses,
                "previous_layer_response": state.previous_layer_response,
            },
        )
# Define a new graph
workflow = StateGraph(HRSystemAgentState)
# Add the node to the graph
workflow.add_node("detect_lang", detect_lang)
workflow.add_edge("__start__", "detect_lang")
workflow.add_edge("detect_lang", "moa_layer_agent")
workflow.add_node("moa_layer_agent", moa_layer_agent)
workflow.add_node("proposer_agent", proposer_agent)
workflow.add_node("aggregator_agent", aggregator_agent)
workflow.add_conditional_edges(
    "moa_layer_agent",
    execute_proposer_agents_or_aggregator_agent,
    ["proposer_agent", "aggregator_agent"],
)
workflow.add_edge("proposer_agent", "moa_layer_agent")
workflow.add_edge("aggregator_agent", "__end__")
# Compile the workflow into an executable graph
graph = workflow.compile()
graph.name = "HR System Agent"  # This defines the custom name in LangSmith

Details about implementation:

1. It was possible to use the prompt of aggregation in each agent from the paper for our use case.
2. To determine the optimal configuration, it was necessary to experiment with variations in the number of agents per layer and the number of layers.
3. As each proposer agent processing the entire context is token expensive, I also experimented with first using the PAE(enhanced version) workers to extract relevant context first and then feeding only the relevant context to each MoA proposer agent. In other words, MoA becomes the final manager agent of PAE. This is more efficient as for the small additional expense of running a PAE relevant context extractor chain we can substantially reduce the irrelevant token consumption per MoA proposer agent. Of course, the Vector similarity search is another alternative for reducing context per proposer agent if your app has too much context.

Evaluation Results(Average scores over 5 runs):

💡Findings:

1. First of all we can see that almost all variations of MoA outperform SingleAgent gpt4o which matches the claims in the paper. The configuration of 2 proposer layers with 6 agents each with Full context added to each proposer using gpt4o-mini gave the best score (+6% vs SingleAgent gpt4o!) while reducing the cost.

2. However the latency for the best configuration is 6x so there is need to choose the appropriate configuration by balancing the trade off between performance and latency based on a usecase or design the UX in a way that intermediate outputs are displayed to the user.

3. We see a clear pattern that performance improves with number of proposer agents with a saturation at 3x6.

4. We see that MoA with PAE improves cost and latency with almost similar performance to Full context MoA so it recommended for practical use in apps. In addition using gemini-flash could reduce cost and latency while maintaining performance but remains to be seen. MoA with PAE can be an indicator of how MoA performs for RAG apps which the paper hasn’t directly addressed!

5. It was unclear why there wasn’t a clear pattern of performance improving with number of proposer agents in MoA with PAE.

3. Multi Agent Debate

• Improving Multi-Agent Debate with Sparse Communication Topology
• Learning to Break: Knowledge-Enhanced Reasoning in Multi-Agent Debate System

The experiment has been conducted by merging ideas from both papers.

What does it do?

Multiple debate agents, after generating their own independent responses, exchange their responses with their neighbor agents and refine their responses based on the responses from their neighbor agents for multiple rounds. A number of rounds can be determined statically or dynamically by having a judge agent judge the responses of debate agents for accuracy, consistency, and consensus and order them to continue with more rounds if needed. Finally, at the end of the last debate round, the responses of all debate agents are sent to a Summarizer agent that aggregates their responses on the assumption that they are all correct and agree with each other. In certain cases, a majority vote (median answer) could be applied to choose the final answer, such as a numerical answer to a Math problem.

Implementing HRSystemAgent

LangGraph Code(4 Debate Agents debating for 3 rounds):

"""This module defines the state representation for the HR System Agent."""
import operator
from dataclasses import dataclass, field
from typing import Annotated, Literal, TypedDict
from pydantic import BaseModel
class Message(BaseModel):  # noqa: D101
    role: Literal["user", "assistant"]
    content: str
@dataclass
class HRSystemAgentState:
    """A state representation for the HR System Agent."""
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str = field(default="ja")
    agent_response: str = field(default="")
    debate_agents_responses: Annotated[list[str], operator.add] = field(
        default_factory=list
    )
    previous_round_responses: str = field(default="")
    current_round_index: int = field(default=0)
class MADAgentState(TypedDict):
    """A state representation for the Multi Agent Debate Agent."""
    user_message: Message
    chat_history: Annotated[list[Message], operator.add]
    user_message_language: str
    neighbor_agents_responses: str | None
    your_previous_answer: str | None
    debate_agent_index: int
"""Define the graph for the HR System Agent.
This module defines a custom graph for the HR System Agent based on Chain of Agents (CoA) architecture.
"""
def detect_lang(state: HRSystemAgentState) -> dict[str, Any]:
    """Detect the language of the user message."""
    language = detect(state.user_message.content.replace("\\n", "")).get("lang")
    return {
        "user_message_language": language,
    }
initial_answer_instructions = """ 
Act as the experienced HR professional that you are for Gaudiy.
Thoroughly review the HR system information provided here: 
{context}
Find the most relevant parts from the above information and use it to answer the user's last message in this chat history:
{chat_history}
"""
debate_instructions = """
There are a few other agents assigned the same task, it's your responsibility to discuss with them and think critically. You can update your previous answer with other agents'
answers or given context as advice, or you can maintain your previous answer. Output your new answer to the user's last message in the chat history below.
Context:
{context}
Other agents' answers:
{neighbor_agents_responses}
Your previous answer:
{your_previous_answer}
Chat History:
{chat_history}
"""
aggregation_instructions_v1 = """
You are an intelligent summarizer agent, tasked with synthesizing the responses of other agents into a concise and comprehensive final answer. Your role is not to
generate original responses, but to condense the information provided by other agents into a succinct summary. You should maintain the writing style and tone of the other agent responses.
Responses from other agents:
{last_round_responses}
"""
async def debate_agent(state: MADAgentState) -> dict[str, list[str]]:
    """Debate Agent"""
    if state["user_message_language"] == "en":
        domain_agent_initial_instructions = domain_agent_initial_instructions_en
        context = context_en
    else:
        domain_agent_initial_instructions = domain_agent_initial_instructions_jp
        context = context_jp
    chat_history = "\\n".join(
        [f"{message.role}: {message.content}" for message in state["chat_history"]]
    )
    chat_history = (
        chat_history
        + f"\\n{state['user_message'].role}: {state['user_message'].content}"
    )
    if "neighbor_agents_responses" in state:
        debate_agent_instructions = debate_instructions.format(
            context=context,
            chat_history=chat_history,
            neighbor_agents_responses=state["neighbor_agents_responses"],
            your_previous_answer=state["your_previous_answer"],
        )
    else:
        debate_agent_instructions = initial_answer_instructions.format(
            context=context,
            chat_history=chat_history,
        )
    system_prompt = f"""
{domain_agent_initial_instructions}
Keep in mind that the current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.
{debate_agent_instructions}
"""
    prompt_messages = [{"role": "system", "content": system_prompt}]
    model = ChatOpenAI(model="gpt-4o-mini")
    response = await model.ainvoke(prompt_messages)
    return {
        "debate_agents_responses": [
            "DebateAgent"
            + str(state["debate_agent_index"])
            + ": "
            + str(response.content)
        ]
    }
async def aggregator_agent(state: HRSystemAgentState) -> dict[str, str]:
    """Aggregator Agent."""
    system_prompt = f"""
    {aggregation_instructions_v1.format(
        last_round_responses=state["previous_round_responses"]
    )}
    """
    prompt_messages = [{"role": "system", "content": system_prompt}]
    model = ChatOpenAI(model="gpt-4o-mini")
    response = await model.ainvoke(prompt_messages)
    return {"agent_response": str(response.content)}
async def mad_round_agent(state: HRSystemAgentState) -> Any:
    """MAD Round Agent."""
    if state.debate_agents_responses:
        previous_round_responses = ""
        for i in range(1, 5): # vary here to change number of debate agents
            if previous_round_responses:
                previous_round_responses += "\\n=====================\\n"
            previous_round_responses += state.debate_agents_responses[-i]
        return {
            "previous_round_responses": previous_round_responses,
            "current_round_index": state.current_round_index + 1,
        }
async def execute_debate_agents_or_aggregator_agent(
    state: HRSystemAgentState,
) -> list[Send] | Send:
    """Decide whether to run the debate agents or the aggregator agent."""
    if state.current_round_index == 0:
        return [
            Send(
                "debate_agent",
                {
                    "user_message": state.user_message,
                    "chat_history": state.chat_history,
                    "user_message_language": state.user_message_language,
                    "debate_agent_index": i + 1,
                },
            )
            for i in range(4) # vary here to change number of debate agents
        ]
    elif state.current_round_index < 3: # vary here to change number of debate rounds
        previous_round_responses_list = state.previous_round_responses.split(
            "\\n=====================\\n"
        )
        to_send: list[Send] = []
        for i in range(len(previous_round_responses_list)):
            your_previous_answer = previous_round_responses_list[i].split(": ", 1)[1]
            debate_agent_index = (
                previous_round_responses_list[i]
                .split(": ", 1)[0]
                .split("DebateAgent", 1)[1]
            )
            other_agents_responses = (
                previous_round_responses_list[:i]
                + previous_round_responses_list[i + 1 :]
            )
            neighbor_agents_responses = ""
            for response in other_agents_responses:
                if (
                    (
                        debate_agent_index in ["1", "4"]
                        and response.split(": ", 1)[0]
                        in [
                            "DebateAgent2",
                            "DebateAgent3",
                        ]
                    )
                    or (
                        debate_agent_index in ["2", "3"]
                        and response.split(": ", 1)[0]
                        in ["DebateAgent1", "DebateAgent4"]
                    )
                ):
                    if neighbor_agents_responses:
                        neighbor_agents_responses += "\\n\\n"
                    neighbor_agents_responses += response
            to_send.append(
                Send(
                    "debate_agent",
                    {
                        "user_message": state.user_message,
                        "chat_history": state.chat_history,
                        "user_message_language": state.user_message_language,
                        "neighbor_agents_responses": neighbor_agents_responses,
                        "your_previous_answer": your_previous_answer,
                        "debate_agent_index": debate_agent_index,
                    },
                )
            )
        return to_send
    else:
        return Send(
            "aggregator_agent",
            {
                "user_message": state.user_message,
                "chat_history": state.chat_history,
                "user_message_language": state.user_message_language,
                "previous_round_responses": state.previous_round_responses,
            },
        )
# Define a new graph
workflow = StateGraph(HRSystemAgentState)
# Add the node to the graph
workflow.add_node("detect_lang", detect_lang)
workflow.add_edge("__start__", "detect_lang")
workflow.add_edge("detect_lang", "mad_round_agent")
workflow.add_node("mad_round_agent", mad_round_agent)
workflow.add_node("debate_agent", debate_agent)
workflow.add_node("aggregator_agent", aggregator_agent)
workflow.add_conditional_edges(
    "mad_round_agent",
    execute_debate_agents_or_aggregator_agent,
    ["debate_agent", "aggregator_agent"],
)
workflow.add_edge("debate_agent", "mad_round_agent")
workflow.add_edge("aggregator_agent", "__end__")
# Compile the workflow into an executable graph
graph = workflow.compile()
graph.name = "HR System Agent"  # This defines the custom name in LangSmith

Details about implementation:

1. Prompts are basically as-is from the paper with minor modifications to incorporate sparse connections between debators and our chat agent use case.
2. Judge agent has been skipped and debate rounds are set statically but could be added trivially.

Evaluation Results:

💡Findings & promising possibilities:

1. We see Multi Agent Debate using gpt4o-mini outperforms SingleAgent gpt4o. The cost though is half that of gpt4o with a latency of 5.5x. A configuration of 4 debate agents and 2 rounds might be more cost and latency effective while also giving better performance than gpt4o.

2. Similar to MoA using a cheaper model like gemini-flash or oss model could further reduce cost and latency. And it may be worth to experiment with only 2 debate agents as well to further reduce cost.

3. Adding a judge agent might further help reduce cost by dynamically controlling the appropriate rounds of debate!

Judge Agent Prompt:

As a judge agent, your primary responsibility is to impartially evaluate the responses of other agents for consistency. Please ensure your judgments are objective, relying solely on the coherence and alignment of the provided answers.

Conclusion

We looked at some multi-agent architectures and realized that all of them outperform SingleAgent for the same LLM, and several of them even outperform SingleAgent GPT-4o at a much lower cost. Increased latency largely remains a problem for them, but we saw that certain promising ways exist to control it.

This was part 2 of a three-part series. In part 3, I will discuss some of the applications we have built internally using these Multi-Agent architectures!

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

Experimenting with Multi Agent Architectures — Can they outperform SingleAgent GPT-4o? was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Shreyas Shinde, AI Engineer of Gaudiy Inc.

Building high-quality LLM applications is no easy task, especially when the latest advancements in AI multi-agent systems are scattered across research papers, blog posts, and experimental projects.

In this article, we cut through the noise and provide a state-of-the-art overview of multi-agent architectures, offering practical insights you can apply to your own AI development.

About Gaudiy’s role in the world of AI Agents

Gaudiy’s AI team vision states that we aim to simulate complex societies through the evolution of generative AI and computing power.

Initially, we focus on breaking down specific tasks and workflows, enabling LLMs to achieve human-equivalent performance. Beyond that, we envision a world where Agents, as individual members, collaborate and coordinate to tackle higher-level challenges and goals related to society and communities.

What is LLM Agents?

The LLM-based system is capable of Context-aware Reasoning on a level of Autonomy.

Context Awareness

• The Agent may be instructed to follow a specific Persona
• The Agent may be provided with Memories and/or instructed to generate new ones and use them appropriately

In terms of LLM Agent, these are achieved by specific techniques such as:

Prompt Engineering ⇒ Effective ways of instructing an LLM to behave (such as adding pairs of expected input/output as examples)

1. Retrieval Augmented Generated ⇒ Provide custom context to the LLM by fetching based on the current input
2. Fine Tuning = Make the LLM learn by updating it’s hyperparameters

Reasoning

• The Agent may generate a plan of several steps to take in order to achieve its goals. Or even generate several plans, act on each, evaluate, and refine them. It may want to divide the task at hand into sub tasks and so on..
• The Agent may choose to use tools like web browsing/data fetching if it thinks it is needed

In terms of LLM Agent, these are achieved by specific techniques such as:

1. Prompt Engineering

• Chain of Thought
• Reflection
• Tree of Thought, etc.（cf. Prompting Guide Techniques)

2. Function Calling

Presenting a list of function descriptions to an LLM and making it decide if it needs to call any function and output the function. These functions could be web api calls or messages to humans, effectively allowing the LLM to interact with its environment.

In a practical sense, these cognitive capabilities are at the level of autonomy. Each level of autonomy has its pros/cons, and an appropriate choice needs to be made by creators

💡As you develop these Single-agent systems, they might grow more complex over time, making them harder to manage and scale. For example, you might run into the following problems:

• agent has too many tools at its disposal and makes poor decisions about which tool to call next
• context grows too complex for a single agent to keep track of
• there is a need for multiple specialization areas in the system (e.g., planner, researcher, math expert, etc.)

To tackle these, you might consider breaking your application into multiple smaller, independent agents and composing them into a Multi-agent system.

Sources:

• https://arxiv.org/abs/2308.11432
• https://arxiv.org/abs/2309.02427
• https://blog.langchain.dev/what-is-a-cognitive-architecture/
• https://langchain-ai.github.io/langgraph/concepts/multi_agent/
• https://lilianweng.github.io/posts/2023-06-23-agent/

Let’s take a small detour to how Anthropic defines Agents, which will not only be a good preview for the multi-agent architectures that will follow but will also help broaden understanding of the state of AI Agents in 2025.

Anthropic came out with it in December 2024 and created quite a buzz on X

Summary of Agents on Anthropic

Anthropic introduces the concept of “agentic systems,” which encompasses both workflows (LLMs and tools arranged in predetermined steps) and agents (LLMs directing their own process without rigid predefinition). Below are the key ideas:

1. Core Building Block: The Augmented LLM

An LLM enhanced with features like retrieval, tools, and memory is at the heart of any agentic system. This empowers the model to generate its search queries, decide which tools to use, and determine what information to store.

2. Workflows: Predefined Code Paths

• Prompt Chaining: Break a task into sequential steps, with each LLM output feeding into the next. Useful for tasks easily split into subtasks (e.g., writing text and then translating it).
• Routing: Classify inputs first, then direct them to different specialized processes or models. Ideal for scenarios where certain types of queries need distinct handling.
• Parallelization: Run multiple subtasks at once or run the same task multiple times to gather diverse outputs. This speeds up execution or increases confidence in the final result.
• Orchestrator-Workers: One LLM breaks down tasks and delegates them to others, then compiles their results. Suitable for problems where you can’t predefine the exact number or nature of subtasks.
• Evaluator-Optimizer: One LLM generates output, another evaluates and offers feedback, and they loop until the quality meets a certain standard. Great for tasks benefiting from iterative refinement.

3. Agents: Autonomous, Flexible Systems

• Agents can operate for many turns without a fixed path, relying on the LLM’s ability to plan, execute, and adapt. They may check back with a human or gather real-world data as they go.
• While powerful for open-ended tasks, agents can accumulate errors over prolonged operations. Hence, thorough testing and guardrails are essential.

4. Putting It All Together

These patterns are not strict recipes. Instead, developers can blend them to suit different use cases, adding complexity only when there is a clear performance benefit.

💡 In conclusion, different players in Academia and/or Industry do not agree on what terms to use for different concepts.

• What one calls an Agent, others may call an Augmented LLM; what one calls Multi-Agents, others may call Agentic Workflows, and purists may consider a system as an Agent only if it is completely Autonomous, and some might just accept that the Agentic behavior is on a level of Autonomy.
• Each player in Academia/Industry seems to be marketing AI Agents from their own perspective to their respective user base for their respective motives.

Now let’s look at how Multi-Agent Systems differ from Single Agent Systems

How are Multi-Agent Systems different?

• The construction of single-agent systems concentrates on formulating their internal mechanisms and interactions with the external environment.
• On the other hand, LLM-MultiAgent systems emphasize diverse agent profiles/personas, inter-agent interactions, and collective decision-making processes. From this perspective, more dynamic and complex tasks can be tackled by the collaboration of multiple autonomous agents, each of which is equipped with unique strategies and behaviors and engaged in communication with one another. (Note that , again these capabilities are on a level of autonomy appropriately chosen by creators)

Agents Profiling methods:

1. Pre-defined explicitly by creators
2. LLM generated
3. Data-derived: generated based on pre-existing datasets

Agents Communication:

1. Paradigms:

• Cooperative → Exchanging info and working towards a common goal
• Debate → Goal is to reach a consensus in a practical amount of duration
• Conflicting/Competitive → Agents work towards their own goals which might conflict with other agent’s goals

2. Structures:

3. May Acquire new capabilities:

1. Feedback from the environment they are operating in.
2. Feedback from other agents.
3. Feedback from humans ⇒ Human in the loop or Human-as-a-Tool
4. Adjust based on new Memory, Self-Evolution (Update its initial instructions), or even generate ad-hoc new agents at runtime!

4. Agents Orchestration Methodologies:

• These methodologies aim to optimize agents' workflows, task assignments tailored to different agents, and communication patterns across agents, such as communication constraints between agents. Effective agent orchestration facilitates harmonious operation among agents, minimizing conflicts and redundancies

These systems are applied and studied in various complex domains:

💡In summary, by varying Agent profiles, How they communicate, & How they acquire new capabilities, we can design appropriate systems or architecture for a use case!

A simple Example:

Sources:

• https://github.com/taichengguo/LLM_MultiAgents_Survey_Papers
• https://medium.com/@pallavisinha12/understanding-llm-based-agents-and-their-multi-agent-architecture-299cf54ebae4
• https://techcommunity.microsoft.com/blog/aiplatformblog/the-future-of-ai-exploring-multi-agent-ai-systems/4226593

What are the properties of Multi-Agent systems/architectures?

1. Single Responsibility Principle/Modularity
2. Breakdown of Complex Tasks
3. Collaboration
4. Scalability
5. Fault Tolerance
6. Parallelism
7. Reusability
8. Need to control Costs
9. Need to control Latency or design appropriate UX

💡Interesting to note that there are similarities with Monolithic vs Modular/MicroServices regarding Software Architecture!

Multi-Agents in Academia & Industry

MA Architectures/Patterns

Below is a categorized list of multi-agent architectures/patterns highlighted in various research papers. They, in some way, try to achieve superiority over single-agent architecture. We will be going through some of them via a deep dive into the paper or an actual experiment.

1. Orchestration Patterns

a. Route to the most appropriate domain expert agent based on the input

b. Decompose the task first and then route to 1 or several appropriate domain agents

c. Chain of Agents

The full context is distributed among several agents, and generations are sequentially passed on as additional context to the next agent from the previous agent. At the end of the worker agents’ chain, a manager agent synthesizes the information into a coherent final output.

d. Long Agent

A leader agent organizes a team to gather information and resolve conflicting information.

e. Dynamically generate Expert Agents based on Evolutionary Algorithms

2. Multiple Generations followed by Selection/Aggregation Patterns

a. Generate multiple candidate codes and verify them using an appropriate method, such as unit testing.

b. Generate answers by multiple Agents and choose best by Majority voting

c. Multiple Agents generate 3 candidate answers each and a final aggregate prediction is made

d. Mixture of Agents： Surpasses GPT4o by using much weaker LLMs in a Multi-layer Generate & Aggregate structure

e. Multi-Agent Debate/Discussion

Multiple agents debate in a certain way and come to a conclusion or consensus about a task.

3. Weak-Strong Agent Combination Patterns

a. A Talker-Reasoner Architecture

Async System 1 thinking(Talker Agent = Fast) + System 2 thinking(Reasoning Agent = Slow) collaborating via shared memory.

Another paper of the same pattern: https://arxiv.org/abs/2305.17390

b. LLM Cascades via Mixture of Thoughts

c. A stronger agent fills out the important details in a template answer generated by a Weaker Agent (this is just my idea. It is similar to https://arxiv.org/abs/2403.03870）

4. System of Agents

a. Internet of Agents

b. AIOS

5. Theory of Mind Based

Multi-Modal Multi-Agent Theory Of Mind

Applications of MA Architectures/Patterns in specific domains

Below, we see some examples of applying the MA architectures in specific domains such as Legal, Social games, Software Engineering, & Machine Learning.

a. Enhance Legal Reasoning with Insights from Multi-Agent Collaboration

b. Multi-LLM Based Agent Debugger via Refinement and Generation Guidance

c. A multi-agent LLM framework for full-pipeline AutoML

d. Weak — Strong Agent Collaboration in Minecraft based on TOM（Theory of Mind）

MA Applications/Systems

Below are some examples of actual MA applications

a. Social Simulations

b. AI Lab

c. Data Analysis Agent inside LinkedIn

d. Interesting Storytelling because Single Agents don’t produce creative stories

MA/SA frameworks

Below is a short list of code or no-code frameworks for Multi-Agents and generically for Single Agents. The list is endless, and new frameworks are being released daily, so these are the main ones.

1. Openai swarm https://cookbook.openai.com/examples/orchestrating_agents
2. Langgraph UI
   https://blog.langchain.dev/langgraph-studio-the-first-agent-ide/
3. Crew AI
   https://github.com/crewAIInc/crewAI
4. LlamaIndex
   https://www.llamaindex.ai/blog/introducing-workflows-beta-a-new-way-to-create-complex-ai-applications-with-llamaindex
5. HF’s SmolAgents https://huggingface.co/docs/smolagents/examples/multiagents
6. https://microsoft.github.io/autogen/stable/
7. https://github.com/microsoft/TinyTroupe
8. https://github.com/geekan/MetaGPT
9. https://docs.langflow.org/
10. https://github.com/awslabs/multi-agent-orchestrator?tab=readme-ov-file
11. https://github.com/lmnr-ai/flow
12. https://flowiseai.com/
13. Gumloop
14. n8n
15. Interframework communication
    https://github.com/langchain-ai/agent-protocol

Personally recommend to start with Langgraph:

Pros:

1. Low entry barrier. Easy to start. UI to debug.
2. Full control and customizable.
3. Modular features so use what you need. Simple concepts like Workflow of Nodes + State ⇒ Graph

Deep Dive Some Papers

Let’s look deeply into some papers to get a feel for how the patterns are used to solve actual problems

Long Agent

https://arxiv.org/abs/2402.11550

Patterns we will come across in this paper

1. Task Decomposition & Iterative Querying
2. Task-based Agent Role Assignment
3. Context distribution
4. Conflict Resolution cycles
5. Aggregation

There are 2 types of Agents

1. Leader Agent x 1
2. Member Agents x n
   a. Entire Context(Single/Multiple Documents) is broken into n chunks, and each member is assigned 1 chunk.
   b. According to the nature of the task, the appropriate role for the member is chosen Ex. QA, Summarizer, Math, Retrieve

Let’s run through an Example:

Query(Multi-Hop query): Which team does the player named 2015 Diamond Head Classic’s MVP play for?

Step 1: Leader Decomposes Complex Queries

The leader breaks down complex queries into simpler sub-questions or sub-tasks

Example: The leader divides the main question into two sub-questions:

“Who won the MVP of the 2015 Diamond Head Classic?”
“Which team did that player play for?”

Step 2: Collaborative Reasoning and Information Gathering

The leader asks the sub-questions to the member agents

Each member agent responds based on the content of its assigned text chunk. Some members may not have the answer in their text chunk

Example:

Member 1 (chunk 1) answers: “Buddy Hield” ⇒ Correct Answer (When the answer exists in a chunk, it is correctly returned)

Member 2 (chunk 2) answers: “Mark Gibson” ⇒ Hallucination (But when the answer doesn’t exist, sometimes the LLM, instead of saying It doesn’t know, will lie)

Member 3 (chunk 3) answers: “No mention in my doc.” (Correctly says don’t know)

Step 3(Key Step): Conflict Resolution

If a member responds to conflict, the leader initiates a conflict resolution process

Members with conflicting responses share their text chunks with each other

Example: The leader detects that members 1 and 2 gave different MVP names and asks them to share their text chunks. This may lead to one or both members correcting their answers.

Step 4: Iterative Querying

The leader can pose additional sub-questions or ask for clarifications if the initial responses are insufficient

This process continues until the leader believes he has enough information to answer the original query

Example: The leader may have a second sub-question, “Which team does Buddy Hield play for?” after determining that member 1 is correct.

Step 5: Answer Deduction

The leader deduces the final answer based on the information gathered through the collaborative process

The leader provides the final answer to the user

Example: The leader determines that Buddy Hield is the MVP and that he plays for the Sacramento Kings. It then provides the final response: “Buddy Hield won the 2015 Diamond Head Classic MVP and plays for the Sacramento Kings.”

Key Points

Contextual Understanding: Each member has a limited view of the text, but by working together, they can extract information from it. The system is designed to accommodate the fact that some members will not have the answers to some questions

Collaboration: The leader and member agents work together, with the leader coordinating the process and the members providing the information from their respective chunks of text

Iterative Process: The system is iterative; it may ask members multiple questions to clarify and refine answers, and members may exchange chunks to resolve differences. This is not a one-shot-question-answer approach

Evaluation

Key Evaluation Metric: Needle In Haystack Plus It is an improvement over the traditional Needle In Haystack

Needle In Haystack

In this setup, a fact or statement of interest (the needle) is placed within a lengthy distracting document (the haystack), and the model is tasked with retrieving this hidden key information. Performance is evaluated by varying the position of the needle within the distracting document and the length of the distracting document itself.

Needle In Haystack Plus

To assess the longtext capabilities more comprehensively, we propose Needle-in-a-Haystack PLUS, which shifts the focus from simple fact retrieval to more challenging single-document/multi-document question answering tasks. In Needle-in-a-Haystack PLUS, the needle represents the document(s) containing the answers, while the haystack comprises distracting documents. The model must locate one or more relevant documents scattered within the haystack and reason the correct answers from them.

Evaluation Dataset:

For multi-document QA, questions require information from two or more documents to reason the final answer. We construct test samples based on HotpotQA, a widely adopted multi-document QA dataset. We select 60 questions from the validation set of HotpotQA that require information from two documents to answer. Each sample contains two needles, and the haystack is still composed of distractor documents. The length of haystack ranges from 1, 000 to 128, 000 tokens with equal intervals, totaling 15 different lengths. The two needles are randomly scattered at the depth of 0%, 33%, 66%, and 100% of the haystack, resulting in 6 combinations: (0%, 33%), (0%, 66%), (0%, 100%), (33%, 66%), (33%, 100%), and (66%, 100%). For specific lengths and needle positions, we randomly select 10 needles to construct 10 test samples for evaluation.

Evaluation Results:

More Red is worse. We can see the loss in the middle phenomenon in GPT4. Using 3.5-turbo for both leader and members gives scores 1.5% better than GPT4.

Conclusions

💡By assigning separate chunks of context to each member, we effectively enhance the system's attention, surpassing that of GPT4. But that introduces hallucinations from members who don’t have relevant context chunks. This is fixed by orchestrating conflict resolution by sharing context among conflicting members (Authors suggest Debate or Reflexion as alternate conflict resolution strategies)

The responsibility of the leader to decompose the query is critical, so it is important to design this agent well. Complex tasks can be accomplished.

Multi-Agent Debate(LLM as a Jury)

Introduction of ChatEval

Introduces a novel framework that enhances text evaluation by employing multiple Large Language Models (LLMs) to debate and assess generated responses.

Code and Demo: https://github.com/thunlp/ChatEval?tab=readme-ov-file#-simple-video-demo

Key Components of ChatEval

1. Debater Agents: Each LLM functions as an individual agent with a unique role or persona, ensuring diverse perspectives during evaluation.
2. Communication Strategies: The framework employs various strategies to facilitate interaction among agents, such as:

• One-By-One: Agents take turns generating responses in a set order. ⇒ Speaking Order Bias
• Simultaneous-Talk: Agents generate responses asynchronously within each iteration.
• Simultaneous-Talk-with-Summarizer: An additional LLM summarizes the discussion at the end of each iteration to inform subsequent responses.

3. Diverse Role Specification: Assigning distinct personas to each agent ensures evaluations from multiple perspectives, enhancing the assessment’s depth and reliability.

Example Scenario

Objective: Evaluate two AI-generated responses to the question: “How can I improve my time management skills?”

Responses:

• Assistant 1: “Improving your time management skills involves setting clear goals, prioritizing tasks, and minimizing distractions.”
• Assistant 2: “Here are some tips to improve your time management: create a schedule, delegate tasks when possible, and take regular breaks to maintain productivity.”

Evaluation Process:

1. Initialization: Three LLM agents are assigned distinct roles:

• General Public: Focuses on the response’s relevance to a general audience.
• Critic: Assesses the clarity and fluency of the responses.
• News Author: Evaluate the consistency and informativeness of the responses.

2. Debate Rounds:

Round 1:

• General Public: Analyzes both responses for general applicability and provides feedback.
• Critic: Reviews the language quality and points out any ambiguities.
• News Author: Checks for factual accuracy and comprehensiveness.

Round 2:

• Agents discuss each other’s evaluations, highlighting strengths and weaknesses in the responses.

Round 3:

• A consensus is reached, with each agent assigning a score to both responses based on the discussion.

3. Final Assessment: The scores are aggregated to determine which response better addresses the question, considering all evaluated aspects.

Advantages of ChatEval

• Enhanced Accuracy: By simulating a human-like evaluation process through multi-agent debates, ChatEval aligns more closely with human assessments.
• Diverse Perspectives: Incorporating agents with varied roles ensures a comprehensive evaluation from multiple angles.
• Scalability: Automating the evaluation process reduces the need for extensive human labor, making it more efficient.

💡But the Simultaneous talk style debate in ChatEval leads to n*n communication problem! Each Agent’s prompt complexity and cost scales with the total number of debaters!

Improved Communication Strategy with lower cost!

Sparse Communication Topology:
Source: https://arxiv.org/abs/2406.11776

We can infer from the above evaluation results that only communicating with your neighbors gives the same or better results over the fully connected MAD topology!

An example scenario in Sparse MAD

Source: https://microsoft.github.io/autogen/dev//user-guide/core-user-guide/design-patterns/multi-agent-debate.html

There are two types of agents in this pattern: solver agents and an aggregator agent. The solver agents are connected in a sparse manner following the technique described in Improving Multi-Agent Debate with Sparse Communication Topology. The solver agents are responsible for solving math problems and exchanging responses with each other. The aggregator agent is responsible for distributing math problems to the solver agents, waiting for their final responses, and aggregating the responses to get the final answer.

The pattern works as follows:

1. The user sends a math problem to the aggregator agent.
2. The aggregator agent distributes the problem to the solver agents.
3. Each solver agent processes the problem and publishes a response to its neighbors.
4. Each solver agent uses the responses from its neighbors to refine its response and publishes a new response.
5. Repeat step 4 for a fixed number of rounds. In the final round, each solver agent publishes a final response.
6. The aggregator agent uses majority voting to aggregate the final responses from all solver agents to get a final answer and publishes the answer.

More references on MAD

1. Financial Sentiment Analysis:

2. One of the latest papers on MAD
https://arxiv.org/abs/2305.19118

What’s Next?

This is part 1 of the 3-part series, and in part 2, I will go through some of the internal experiments we have done involving some of these multi-agent architecture papers!

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

Considerations on Multi Agents(Part1)- A Comprehensive Survey was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by seya, LLM App Engineer of Gaudiy Inc.

As the name suggests, LLM-as-a-Judge is a technique that uses LLMs for various types of evaluations. LLMs are excellent at interpreting unstructured data, allowing them to make judgments or assign scores to tasks that are typically difficult to assess with conventional programming.

In this article, I’ll introduce my approach and tips for creating an LLM-as-a-Judge system to confirm improvements and detect regressions in LLM-based systems.

Why Use LLMs for Judging?

First, let’s clarify the purpose of why we want to evaluate using LLMs.

The primary goal is to detect whether an LLM system has improved or regressed. LLM-based features can change due to various factors:

• Modifying prompts
• Changing input methods
• Adjusting LLM API request parameters
• Switching models
• Data changes

With so many variables, it becomes time-consuming to manually check for improvements or regressions each time we make changes. Therefore, we want an automated system that can inspect and judge, similar to writing unit tests in programming.

This is the overarching goal we want to achieve.

*Note: While LLM-as-a-Judge can be used as a feature within a product, this article focuses solely on using it for internal quality improvement.

While LLMs are not always the best tool for automated evaluation, for certain evaluation aspects that are difficult or impossible to judge with conventional programming, we want to use LLM-as-a-Judge.

When incorporating LLMs into products, the criteria for evaluating their behavior tend to be much fuzzier compared to testing in traditional system development.

For example, in our company, when developing an AI that mimics a specific person, we have evaluation criterias such as:

• Showing empathy, encouragement, or a listening attitude
• Responding in casual language, not using formal speech

These are aspects that are challenging to evaluate with pure programming. This is where LLMs come into play.

However, as mentioned earlier, “LLMs are not always the best tool.” If an evaluation aspect can be judged with conventional programming, it’s generally better to use that approach. This is because LLM-as-a-Judge is slow and expensive. 💸 (Although recent models like GPT-4o mini and Gemini Flash are becoming less so)

For instance, if you wanted to confirm whether the response is in the same language as the user’s query (English for English, Japanese for Japanese), you could use LLMs to judge this, but there are also dedicated language detection libraries available.

GitHub - LlmKira/fast-langdetect: ⚡️ 80x faster language detection with Fasttext | Split text by language for TTS

So for such evaluation aspects, we use functions like this to generate scores:

from fast_langdetect import detect
from langsmith import EvaluationResult
from langsmith.schemas import Example, Run
def answer_with_language_same_as_query(run: Run, example: Example) -> EvaluationResult:
 """Check if the agent reply is in query language."""
 query = example.inputs["query"].replace("\n", "")
 query_language = detect(query)
agent_reply = get_output(run, "output").replace("\n", "")
 agent_reply_language = detect(agent_reply)
score = query_language.get("lang") == agent_reply_language.get("lang")
 return EvaluationResult(key="ANSWER_WITH_SAME_LANGUAGE_AS_QUERY", score=score)

There are also evaluation aspects that are not suitable for either LLMs or programming. Clear examples include subjective elements like “conversation interestingness” or specialized areas requiring practical experience, such as medical or legal consultations. If you try to evaluate these aspects, LLMs might provide plausible reasons and scores, but the correlation with real human evaluations is likely to be relatively low.

I’ll explain “how to create LLM-as-a-Judge” next, but it takes more than just a quick effort to create this. Therefore, when actually creating it, it’s important to question, “Is LLM really the optimal choice for this evaluation?”

Now that we’ve aligned on the purpose, let’s look at how to create it!

How to Create LLM-as-a-Judge: Start with Annotation

When you actually start creating LLM-as-a-Judge, you might be tempted to first consider the evaluation criteria… but I believe it’s better to start by diligently creating data and performing annotations.

This is because you’ll need datasets for later evaluation anyway, and evaluation criteria often only become clear after you’ve actually done some evaluations. Of course, you should have a rough verbal description based on the functional requirements and hypotheses you have, but at this stage, I think it’s important not to decide too rigidly and to actually try evaluating.

The overall picture, including LLM-as-a-Judge, will eventually look something like this (of course, we won’t go as far as evaluating the evaluation of the evaluation, as that would be excessive):

So, we start by creating a dataset to run the LLM system we want to evaluate, actually running it, and then annotating the results.

Our company uses LangSmith as a platform for evaluation and experiment management, so we ultimately input the data there and use the SDK’s evaluation functions. This is convenient for seeing changes over time when scoring evaluations.

By the way, we’ve released a library as OSS to perform evaluations in LangSmith, so feel free to use it if you’d like.

GitHub - gaudiy/langsmith-evaluation-helper: Helper library for LangSmith that provides an interface to run evaluations by simply writing config files.

As you continue annotating, you’ll likely find yourself able to verbalize evaluation aspects, and these evaluations can later be used to evaluate the evaluation itself.

Request Annotations from Domain Experts

After talking about just diving into evaluation, I don’t mean to dampen the enthusiasm, but it’s important to ask people who can evaluate more accurately — domain experts for the product you’re trying to build. (For the very first PoC, it’s fine if someone from the development team does it for speed, but…)

As a personal failure story, I recently started by annotating the responses of an in-house agent myself. However, my evaluations were quite lenient, and when I later compared annotations with internal people (domain experts) who had the domain knowledge that the agent should assume, many of the items I had judged as OK were determined to be incorrect.

If you looked closely at the LangSmith reference screenshot I just posted and thought, “The score has gone down!”, it’s because I’ve adjusted to a stricter evaluation.

If I had created an LLM-as-a-Judge based on my annotations, the scores might have been unfairly high, unable to detect problems and thus becoming useless, or worse, giving an inaccurate sense of security that could have been harmful.

These evaluation criteria and the resulting scoring are what I consider to be guidelines for how to nurture any LLM product. So I think it’s an unexpectedly important decision.

If you’re only creating something that’s “just for reference,” it won’t lead to the experience of being able to confidently make modifications with reduced effort. So it’s important to find people who can do high-quality annotations, collaborate with them, and improve over time.

Don’t Over-Engineer Evaluation Criteria Too Early

As I mentioned when discussing the sense of purpose at the beginning, “similar to writing unit tests in programming” is important. The evaluation we’re doing with LLM-as-a-Judge is ultimately to confirm whether the LLM system is working as expected or better within the specified parameters.

In other words, even if the accuracy of the evaluation itself is extremely high and you can achieve a high score for the LLM system accordingly, it might still be poorly received by users or contribute nothing to the growth of the product.

In such situations, you might realize that the feature is no longer necessary, or that a different evaluation aspect is important, potentially rendering the data you’ve been diligently creating internally worthless. So investing too much in this automated evaluation system by deciding on evaluation criteria too early could be a waste of time.

Therefore, I think it’s better to start with quality standards that won’t change and must be maintained, and then gradually tune the evaluation items and criteria based on user tests and actual usage after release.

To bring in something I wrote in a different article a while ago, it’s easier to imagine if you think of it as climbing up the categories defined in the Kano model.

• Start with basic quality (the kind of quality protected by guardrails)
• Then one-dimensional quality
• Finally, attractive quality

source: KANO model

Personally, while I believe that evaluation is most important in creating LLM products, I think it’s crucial not to become too dogmatic and to gauge the level of investment according to the situation.

Evaluating the Evaluation

Finally, we create an evaluation dataset for the evaluation prompts and check their accuracy, so I’ll introduce that process.

Create Evaluation Prompts

First, we need evaluation prompts for the target we’re evaluating. Let’s create them.
We’ll do proper tuning after creating the evaluation dataset, so for now, it’s fine as long as it doesn’t look too strange when checked visually.

As an example, here’s a simple one I’ve written:

Key points:

• Have it return reasons — This makes it easier to judge when evaluating the evaluation later
• Temperature is 0 — If the evaluation changes between different experiments, it becomes difficult to track what caused the score to go up or down

class EvalResult(BaseModel): # noqa: D101
 score: str = Field(description="score (0 or 1)")
 score_reason_jp: str = Field(description="reason for the score in japanese")
 score_reason_en: str = Field(description="reason for the score in english")
 
def judge_correctness_with_llm(query: str, agent_reply: str, ground_truth_agent_reply: str) -> EvalResult: # noqa: D103
 llm_as_judge = ChatOpenAI(model="gpt-4o-mini", temperature=0).with_structured_output(EvalResult, method="json_mode")
 llm_judge_result: EvalResult = llm_as_judge.invoke(
 f"""Your task is to give a score(0 or 1) for the below Agent reply based on whether the given agent reply has same information as ground truth agent reply.
Query:
{query}
Ground truth Agent reply:
{ground_truth_agent_reply}
Agent reply:
{agent_reply}
# Criteria
- If the information in the agent reply has nuance that's opposite to ground truth agent reply, give score 0, else give score 1.
- If the agent reply is missing some information from the ground truth agent reply, but the rest of the information is correct, give score 1.
- not all the information have to match exactly.
# Output format
You must respond in json format with the following keys:
- score_reason_jp: reason for the score in japanese
- score_reason_en: reason for the score in english
- score: score (0 or 1)"""
return llm_judge_result

By the way, there are decision points in “how to assign scores,” such as whether to output 0 or 1, or to use a score with more gradation.

Personally, I prefer 0 or 1 because it’s easier to understand and judge, but this too depends on the “evaluation aspect.”
For example, in a recent case, I wanted to determine “how much of the content that should be included in the Ground Truth answer is present in the agent’s response.” A 0 or 1 evaluation would have been either too strict or too lenient for this evaluation aspect, so I changed to a method with gradation by generating a list of “items that should be included” and determining what percentage of them were present.

In this way, how to evaluate and assign scores varies depending on the evaluation aspect.

Also, besides assigning scores, there are various methods like pairwise comparison (preparing two answers and choosing which is better) or taking the average of scores from multiple models.

Create an Evaluation Dataset for the Evaluation

Next, for creating the evaluation dataset, the annotations we’ve been doing come in handy.

In my recent workflow, I first run the system to be evaluated in LangSmith, then pick out items that disagree with human evaluation and compile them into a separate dataset.

This kind of data is optimal for checking false negatives/false positives.
When tuning LLM-as-a-Judge, the goal is basically to align with human judgment (annotation).

What’s useful when aligning is the items that differ from human judgment, and the main focus during tuning is to eliminate “items judged as correct when they’re actually incorrect — false positives” and “items judged as incorrect when they’re actually correct — false negatives.”

Of course, items that already match in judgment are also useful as data for regression checks.

By the way, from my experience so far, I tend to focus more on false positives, i.e., detecting cases that are “not good but given a high score.”
While having something good judged as bad increases wasted checking time, it’s not a big problem if the number of cases isn’t too high.

On the other hand, overlooking problems due to false positives has a greater impact on quality degradation, so prioritizing the expansion of datasets to check for these becomes more important.

Keeping these aspects in mind, we accumulate data. There’s no specific correct amount, but I think about 30–50 cases are sufficient for initial verification.

Expanding the Evaluation Dataset Using Synthetic Data

Using synthetic data, that is, having LLMs expand the test dataset, is also effective in increasing confidence.

Above all, the previous method limits the range of inputs to the initially created dataset. Therefore, it’s useful to ensure a range of tests with inputs that are similar in aspect but different in content.

However, it’s important not to add too many overly simple examples. If you prepare test data that’s the exact opposite of the ideal answer to test for false positives, it might result in data that’s too obvious like “Of course it would judge this as incorrect!” Such data doesn’t make for a good test. In fact, if too much of this data accumulates, it might give a false sense of security like “The accuracy rate is over 90%, so we’re safe.”

Of course, it’s effective to gradually test difficulties that you want to overcome, like “Surely it should handle this much…” In such cases, it’s good to add tags for difficulty levels so they can be executed individually.

As a reference, in a recent example of creating synthetic data, I asked GPT-4o to create slightly incorrect data like this. It can create good examples that aren’t obviously wrong at first glance but are incorrect upon careful reading.

def generate_slightly_wrong_output(reference_answer: str) -> str:
 """Generate a slightly wrong output."""
llm = ChatOpenAI(model="gpt-4o", temperature=0)
 output = llm.invoke(
 f"""Your task is to generate a wrong answer text based on the reference answer.
this output is going to be used to test evaluator's ability to judge the correctness of the information.
reference answer:
{reference_answer}
background information of the agent reply:
{document}
"""
 )
 return str(output.content)

Write an Evaluation Program for the Evaluation

I’m simply checking if it matches the reference:

from langsmith import EvaluationResult
from langsmith.schemas import Example, Run
def check_actual_score_matches_reference_score(run: Run, example: Example) -> EvaluationResult:
 """Evaluate if the output score of the run matches the reference score."""
 actual_score = run.outputs.get("output") if run.outputs else None
 reference_score = example.outputs.get("output") if example.outputs else None
if actual_score is None or reference_score is None:
 raise ValueError("actual_score or reference_score is None")
score = int(actual_score) == int(reference_score)
return EvaluationResult(key="matches_reference", score=score)

Tuning!

Once all the preparation is done, all that’s left is to keep tweaking the LLM system until the score improves.

From my experience, recent LLMs are cheap and smart, so this tuning itself doesn’t seem to take much time. Rather, the key point that takes time is “how to effectively verbalize the evaluation criteria.” So it might feel more like criteria tuning than prompt tuning.

Conclusion

That’s how to create LLM-as-a-Judge.

It takes a bit of effort, but when it successfully detects regressions, you feel good thinking, “Oh… it’s smart…”

In practice, LLM systems often improve in one aspect while deteriorating in another as you tweak them, so this approach often comes in handy.

Let’s make good use of it!

References

• Evaluating the Effectiveness of LLM-Evaluators (aka LLM-as-Judge)
• Who Validates the Validators? Aligning LLM-Assisted Evaluation of LLM Outputs with Human Preferences

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

LLM-as-a-Judge: A Technique for Evaluation Using LLMs was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Shreyas Shinde, AI Engineer of Gaudiy Inc.

Introduction

In the rapidly evolving field of artificial intelligence, ensuring that AI agents can provide relevant and accurate responses is crucial. One of the key challenges in AI conversations is handling inappropriate or ambiguous user queries. This blog explores the importance of query contextualization using chat history to improve agents’ memory retrieval and overall conversation quality. It also goes deep into how to systematically tune a prompt to achieve query contextualization.

At Gaudiy, we provide “Gaudiy Fanlink,” a community platform for multiple IPs (intellectual properties). It has a Beta feature that allows users to chat with AI avatars of specific characters (such as actual celebrities or anime characters) in collaboration with specific IP holders. We use vector cosine similarity search to retrieve the Character AI’s memories (base information about the character and summaries of chat history with a user) relevant to the user’s message query from a Vector Store (using Vertex AI) and then add them to the LLM’s prompt to initiate a reply generation process. As the user’s message doesn’t always contain enough context to retrieve relevant memories it becomes essential to first contextualize the user’s message using recent chat history and then use it as a query.

Problem Example

Imagine a conversation between a User and Agent below:

User: Have you read any new books recently?
Agent: Yes, I recently read ‘Harry Potter’.
User: That's great. Which volume are you reading?
Agent: I've started reading the first volume. It's very interesting.
User: Nice!

When we directly use the user’s last message Nice! as a query for the vector db we would get results which would not be related to what that particular Nice! was about. It would miss out on getting information about key topics and nuance in the current conversation i.e Harry Potter, books, first volume & interesting . If our Agent has memories which are relevant to these topics within the context of Nice! , it would miss out on the chance to utilize them to have a much richer & personalized conversation.

Literature Survey

We investigated 3 different possible approaches:

1.

Prompt an LLM to rewrite the user’s message based on the chat history such that it can be understood standalone. This has been documented in LangChain’s docs.

We quickly found out that although the simple prompt provided by LangChain performs great when used with a powerful but expensive LLM like gpt4o, there were several issues when used with a cost effective LLM like gpt3.5, gemini-1.0, gemini-1.5-flash, claude-sonnet-3 depending upon the model chosen.

Some Issues faced:

1. Replies to the user’s last message instead of creating a contextualized message
2. Only partially contextualizes the message
3. Adds hallucination in the contextualized message

The choice for this LLM needed to be a cost effective one for it be used in our product. We found gemini-1.5-flash to be the most promising among the cheaper models.

2.

To counter the issue of incomplete contextualization mentioned in 1, instead of only 1 generate 3 standalone queries(statement or question) from user message and chat history and retrieve per query.

But this approach has the additional overhead of post-retrieval reranking and filtering the set of memories retrieved by 3 queries. We could use an approach like reciprocal rank fusion or llm filter to do so. We would also need some prompt engineering to improve the relevance of the 3 queries and to avoid adding hallucinations.

3.

HypotheticalDocumentEmbedding: This approach first generates an Agent reply without using any memories and then uses the reply to retrieve memories and then regenerate the reply. As our reply generation process is multi-step this approach leads to doubling our cost so it is not very cost effective.

💡 Decision:

We chose to go ahead with the first approach above as it was judged that by prompt tuning we would be able to recognize and resolve issues and use it in our product cost effectively.

Prompt Tuning for Query Contextualization

In order to improve the simple prompt from LangChain in a structured way we took below steps:

1. Generate a synthetic dataset of various types of conversations between a user and agent for evaluating the prompt against using chatgpt4. Add ground truth contextualized queries per example. These are the so called ideal queries which would lead to the retrieval of most relevant memories Upload the dataset to LangSmith.
2. Write an evaluator function using LangSmith which calculates the cosine similarity distance between the embeddings of the generated contextualized query and the ground truth contextualized query. Shorter distance would mean the prompt is producing a contextualized query for an example that is similar to the ground truth and leads to similar relevant memories retrieval as would with the ground truth query. We used the langsmith helper library we recently released as OSS to help us with this evaluation.
3. Execute the query contextualiser llm per example and evaluate the generated contextualized query. Calculate average score across categories of examples. Identify troublesome examples and update the prompt in an attempt to resolve them. Rerun the evaluation to confirm if they are resolved. This cycle of evaluation not only gave us a direction for improving the prompt on specific category of examples but as we were also computing an average evaluation score during each iteration it also made sure that the prompt was also performing well on the overall dataset.
4. Run it against actual production queries. Add few shot examples and set the LLM’s temperature to 0 for consistent outputs as it reduces the LLM’s creativity.

1. Generating Synthetic Dataset using chatgpt4

By chatting with chatgpt4 we created a high quality synthetic dataset of 75 conversations by varying the types of chat_history and user_last_message. It is important to consider the example types deeply for better coverage of actual scenarios.

We also generated ground truth contextualized queries for evaluation.

Example:

"chat_history": "L: The SPK has valuable intel.
        Shreyas: Should we contact them?
        L: Yes, but cautiously.
        Shreyas: What's the best approach?
        L: A secure line.",
"user_last_message": "Should I set up the call?",
"ground_truth_standalone_message": "Should I set up a secure call to contact the SPK for their valuable intel?",
"example_type": "Multiple turns, question, context from multiple messages"

Uploaded to LangSmith:

What was the reason for Human in the loop style of synthetic dataset generation using chatgpt4?

💡 There were clear merits of doing this task in a chat way rather than one prompt and script. Chat way gave more control over the dataset being generated (I first asked it to generate 10 examples with the given prompt and then verify what mistakes it was doing). As we can see how it is failing to follow the instructions we can immediately ask it to correct it, verify again and then produce more examples. We decide the conditions of how we want to vary the examples (in my case I wanted example types like user_last_message is statement or question, chat_history has single message or multiple messages etc). It also allows us to flexibly make changes to the schema mid-way without having to be perfect about it in the first go in case it is needed to change it. As it has a long context window all the history is being used to keep following all instructions in previous messages. Chatgpt easily outputs files in csv, json formats. CSV file can be directly uploaded from LangSmith’s UI to create a new dataset! It also helped me finally translate the file into Japanese because first I had to verify the data in English. This time we only needed to create 75 examples. In case we need 1000 examples the idea will be that after performing this chat and we can satisfactorily and consistently produce examples we ask it to summarize the conversation and create a prompt. And then use that prompt by script and api to generate 1000 examples.

2. Evaluator function using LangSmith

We used the LangSmith’s off the shelf embedding_distance evaluator

from typing import Any

from langchain_google_vertexai import VertexAIEmbeddings
from langsmith.evaluation import LangChainStringEvaluator
from langsmith.evaluation.integrations import _langchain
from langsmith.schemas import Example, Run

from gaudiy_python_kit.ai_testing.tests.utils import get_child_run, get_output


def prepare_data_for_similarity_eval_single_contextualized_query_and_ground_truth(  # noqa: D103
    run: Run, example: Example | None
) -> _langchain.SingleEvaluatorInput:
    prompt_run = get_child_run(run, "prompt")
    relevant_memories_run = get_child_run(prompt_run, "get_relevant_memories")
    pre_retrieval_query_translation_run = get_child_run(relevant_memories_run, "_pre_retrieval_query_translation")
    contextualize_query_run = get_child_run(pre_retrieval_query_translation_run, "contextualize_query")
    return {
        "prediction": get_output(contextualize_query_run, "output")[0],
        "reference": get_output(example, "ground_truth_standalone_message"),
        "input": "",  # added because mypy complains about missing key
    }


evaluators = [
    LangChainStringEvaluator(
        "embedding_distance",
        config={
            "embeddings": VertexAIEmbeddings(model_name="textembedding-gecko-multilingual@latest"),
            "distance_metric": "cosine",
        },
        prepare_data=prepare_data_for_similarity_eval_single_contextualized_query_and_ground_truth,
    ),
]
summary_evaluators: list[Any] = []

3. Evaluating the query contextualizer’s prompt on the dataset and prompt tuning

1. First of all in order to create a benchmark we ran the evaluator without query contextualization. The score was 0.11 (average of embedding distances between user_last_message and ground_truth_standalone_message).
2. Next we ran evaluation on following configuration:
   a. v1 of the prompt (langchain’s simple prompt)
   b. model: gemini-1.5-flash
   The score was 0.08. As the average distance reduced we could conclude that query contextualization makes the user_last_message more cosine similar to the ground_truth query.
3. We then checked the type of examples which had largest embedding distance as those were the clear targets for prompt tuning. After 22 such prompt revisions we finally reached a score of 0.04 and a good enough prompt that performed well on all the different types of examples. We made sure to version the prompts in code in order to be able to back track and compare variations effectively. Again chatgpt4 helped here in improving the prompts by just giving it the current prompt and a list of examples where the prompt was performing poorly!

💡 We can see from the above graph that we keep tuning the prompt in a manner that the distance keeps reducing!

4. Evaluating against actual production chat log

Running only against synthetic dataset is not enough to confirm the prompt works well in actual practice. It is because synthetic datasets suffer with the problem of textual diversity even if we have made sure to create a variety of examples. For better confidence we tested the prompt against a few production queries and chat history and quickly found some improvements that needed to be done to the prompt. We also found that adding few shot examples to the prompt and setting the LLM’s temperature to 0 leads to much more consistent results. As gpt4o-mini was recently released we also did comparative evaluation using it and found it performed better so we switched to using it.

Results

1. Some actual outputs (in Japanese)

#Chat History:
shreyas: 最近何か新しい本を読みましたか？
AI: はい、最近「ハリーポッター」を読みました。
shreyas: それは素晴らしいですね。どの巻を読んでいますか？
AI: 最初の巻を読み始めました。とても面白いです。
#Last user message:
いいね！
#Contextualized query:
最近「ハリーポッター」の最初の巻を読み始めたことに対して、いいね！と言います。

#Chat History:
shreyas: ピカチュウって、ほんとにあるのかな？
#Last user message:
あるよ！
#Contextualized query:
ピカチュウは実在するキャラクターで、ポケットモンスターの一部として知られていますよ！

#Chat History:
shreyas: 明日、アイドルフェスティバルでライブ配信があるんだ
#Last user message:
それって何ですか？
#Contextualized query:
明日行われるアイドルフェスティバルについて詳しく教えてもらえますか？

2. Prompt and model config

Model: gpt-4o-mini
Temperature: 0
Few shots: Yes (ref: using LangChain)
Prompt:

Examine the provided chat history and the last user message to determine its clarity as a complete query or statement.
The chat history is formatted as a continuous string where each message is on a new line and includes the sender's name, while the last user message is given separately without a username identifier.
If the last message is already a clear, complete query or statement that can be understood independently, return it as is.
If it lacks sufficient clarity or requires contextual details to be fully understood on its own, carefully examine all previous interactions in chat history to gather every relevant piece of information.
If the last message implicitly or explicitly references any aspect discussed earlier in the chat history that is essential for understanding then Rewrite the last message to include all necessary details discussed earlier, ensuring it clearly & completely integrates those pertinent details to stand alone confidently and directly connects the context from chat history to the user's statement.
The rewritten message in japanese should maintain the original message's intent, either as a question or a statement, without explaining or answering itself.

Note: update the prompt if you want output in English!

Conclusion

In conclusion, the approach of contextualizing user queries using chat history has proven to be a critical enhancement for improving memory retrieval in AI conversational agents. By systematically tuning prompts and leveraging synthetic datasets for evaluation, we were able to significantly reduce the embedding distance between the original and contextualized queries. This improvement ensures that AI agents can retrieve more relevant memories, leading to richer, more personalized conversations. The iterative process of prompt tuning, coupled with real-world testing, has resulted in a robust system capable of handling diverse conversation scenarios effectively. This method not only enhances the user experience but also optimizes the performance of AI agents in practical applications. The success of this approach underscores the importance of context in AI conversations and demonstrates the value of carefully designed prompt engineering and evaluation strategies in developing more capable AI systems.

Gaudiy will continue to advance research and development in AI agents.

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

Context Matters: Using Chat History to Retrieve More Relevant Memories for AI Agent was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.

Posted by Wataru Namiki, Software Engineer of Gaudiy Inc.

Over the past 1–2 years, the excitement in the generative AI / LLM field has accelerated significantly, and many people are focusing on how to utilize it to provide new value.

Our company, Gaudiy, saw potential in this field relatively early and has been actively taking on challenges. As we develop LLM products, we’re refining our processes to improve daily based on the knowledge accumulated from addressing challenges that arise.

This time, I’d like to focus on the process of prompt tuning, which is inseparable from development in this field. I’ll introduce how we’re improving efficiency and resolving common issues, along with some use cases.

Note: This article is the second part of a two-part series, covering the “Process Implementation” section. I hope you’ll also read the first part by seya, which covers the “Tech Selection/OSS” section.

Part 1: Tech Stack Selection — How We Used LangSmith to Streamline Evaluation and Experimentation…

What we build — Fanlink

At Gaudiy, we provide “Fanlink,” a community platform for multiple IPs (intellectual properties).

Although it’s still in beta and we can’t share many details, we’re currently developing and iterating on a feature that allows users to chat with AI avatars of specific characters (such as actual celebrities or anime characters) on Fanlink, in collaboration with specific IP holders. We call this feature “AI Talk.”

For this feature, we use LLM to generate AI messages. LLM is an essential component, and it’s not an exaggeration to say that it greatly influences the quality of the feature.

Let me briefly introduce the current AI Talk feature:

Users can send messages to AI avatars (hereafter referred to as character AI) that behave like specific characters on Fanlink. The character AI retains information about its unique characteristics and memories, and is always prompted to behave appropriately in response to user messages and the flow of conversation.

In addition to text messages, we also offer supplementary experiences such as giving digital item gifts. We’re also considering features that allow multimodal inputs like voice in the future.

Prompts include not only replies to user messages but also various scenarios such as initiating conversation topics, self-introductions, and ending conversation flows. We use different prompts in our logic depending on various contexts. This allows us to create a mechanism that can behave more naturally from the start to the end of a conversation segment, or seamlessly transition to the next topic.

Next, let me briefly introduce the architecture:

When a message is sent from a user, we stack a task in the reply generation queue (using Cloud Pub/Sub). When processing begins, we retrieve information from a DB that stores character features based on a unique ID representing which character AI’s information it is.
Note: The character AI itself is structured so that multiple individuals can be built separately, and we’re currently operating and verifying two character AIs.

Furthermore, we use vector search to extract information close to the user’s message query from pre-stored base character AI memory data and summaries of past conversations between users and character AI in the Vector Store (using Vertex AI) .

Then, based on context information such as the start of the conversation or how many times the conversation has occurred, we determine which prompt template to use. We extract only the necessary inputs for the prompt from various data obtained above, such as character information and related memories, and generate messages by invoking LLM using the LangChain interface.

This approach allows us to realize the behavior of multiple character AIs with a single prompt template.

While very simple tasks are composed of single prompts, there are also patterns where we chain-execute prompts in multiple steps using LangGraph for more complex context recognition and cost reduction (we’ll omit the detailed technical content here).

The execution of such LLMs can be traced and viewed on the GUI provided by LangSmith from LangChain, making it very easy to debug when hallucinations occur.

Challenges Faced in Prompt Tuning Tasks

As mentioned earlier, multiple prompts are used to realize the character AI talk feature, and the importance of prompt tuning to improve the accuracy of each one was very high.

However, we hadn’t established an efficient way of prompt tuning, and we were facing various challenges such as ballooning time-costs due to ad-hoc approaches.

The biggest challenge was that the common language between the engineers in charge of tuning and the stakeholders (our Biz members) was only qualitative results based on the final behavior of the application. This resulted in large costs and losses for engineers to break down and analyze abstract requirements to the implementation level.

Previously, after a set of features was fully implemented, we operated with a system where our Biz members, who were well-versed in the AI’s characteristics and in close coordination with the IP, essentially became the stakeholders and were asked to check the quality.
The result of the quality check was a qualitative list of items from an infinite number of perspectives, such as whether it met the characteristics of the base character AI, including speech patterns.

While the quality check results were correct, the burden on engineers to interpret each of these qualitative perspectives and implement them into prompts (sometimes prioritizing them) while meeting all of them was very high.

In this work, where the completion criteria were difficult to see, there were times when the work itself naturally disappeared or development proceeded without satisfactory tuning. In the end, there were even unfounded conclusions such as the quality had decreased because we changed to a cheaper LLM model.

When doing prompt tuning, problems like “if you stand this up, that falls down” often occur. The communication cost of having to ask questions each time about which requirements ultimately needed to be realized as must-haves was also significant.

We also used LangSmith’s Playground to check and tune prompt behavior, but in this case, there was no history left, and each trial had to be done manually, which was very costly.

This approach would also be very costly when comparing between LLM models to determine which is optimal for the prompt in question.

Improving Efficiency in Prompt Tuning

Planning Phase

Although an AI team had been formed and we were gaining more members familiar with this field, the AI team was mainly focused on R&D and couldn’t approach these messy process issues. As a member of the product development team, I felt there was a big issue here, so I thought it would be good if we could at least achieve the following to resolve this:

• Create a dataset in some kind of document with multiple expected inputs and examples of correct answers as expected values
• Prepare an environment that can automatically execute LLM and generate execution results based on the prepared dataset
• Prepare persistent data that can serve as a common language for tuning personnel (engineers) and stakeholders (Biz) to discuss improvements based on execution results
• Easily switch to any LLM model and compare results

Implementation Phase

First, we decided to create a dataset of multiple expected inputs and correct cases in a Google Spreadsheet.
The reason for this choice was that it matched the characteristics of spreadsheets, which allow multiple people to edit in real-time online and easily reflect changes in common definitions to the entire dataset using sheet/cell references.

We structured the spreadsheet as follows:

• A common definition sheet that duplicates DB data to make it easier to tune the character AI’s feature information itself
• Sheets (multiple) for each prompt, summarizing the necessary inputs and outputs for each case in one row, created for as many patterns as needed

We set up select boxes on the sheets for each prompt to set which character AI the tuning case is for, and used formulas in cell values to auto-fill dynamically from the common definition sheet side.

We asked Biz members to create other necessary user message inputs and expected output values for each case to cover all patterns.

Next, for the efficiency of LLM execution and extraction of execution results, we had almost no technical knowledge in this area and were about to start digging when the AI team volunteered to collaborate. After re-syncing the sense of issues and the world view we wanted to realize, they proposed a mechanism to batch execute for each dataset and persist execution verification results on LangSmith using LangSmith’s Datasets & Testing feature!

Since the API for effectively using LangSmith’s Datasets & Testing feature is published by LangChain, we were able to realize our requirements by wrapping these for our use case.

*This Helper is currently also published as an OSS library. For more details, please read the repository mentioned below or seya’s article introduced at the beginning.

GitHub - gaudiy/langsmith-evaluation-helper: Helper library for LangSmith that provides an interface to run evaluations by simply writing config files.

From here, I’ll briefly illustrate the overall realization flow and show part of the implementation for efficiently conducting prompt Evaluation along with how to use the Helper (please note that it may differ slightly from the current Helper OSS interface).

First, we implement and execute a Python script (dataset.py) to read information from Google Spreadsheet, convert it to the format of LangSmith’s Datasets Example, and register it.

For reading, we used the Google Sheet API provided within GCP, which Gaudiy usually uses as infrastructure.
To use the Google Sheet API from the script, we realized it by granting viewing permissions to the previously created service account in the sharing settings on the spreadsheet side (note that the way of assigning permissions is different from IAM).

We perform necessary mapping and formatting processing on the read dictionary data, and call the Helper to register/overwrite to any Dataset Example on LangSmith.

# LangSmith Dataset Name
LANGSMITH_DATASET_NAME = "Test Dataset"
# Google Spreadsheet ID
SPREADSHEET_ID = "hogehoge"
# Google Spreadsheet name
SHEET_RANGE = "test_prompt_sheet"
# path to SA json for Google Sheet API
SERVICE_ACCOUNT_KEY_PATH = os.getenv("SERVICE_ACCOUNT_KEY_PATH")

async def dataset() -> None:
        # fetch Spreadsheet
    values = get_sheet_values(SHEET_RANGE)

    # map and create input data
    header_row = values[0]
    header_columns_count = len(header_row)
    data_rows = values[1:]

    dict_list: list[dict[str, Any]] = []

    for row in data_rows:
        name = row_dict.get("name", "")
        bio = row_dict.get("bio", "")
        attitude = row_dict.get("attitude", "")
        last_message = row_dict.get("last_message", "")
        output = row_dict.get("output", "")

        dict_list.append({
            **TestPromptInput(
                name=name,
                bio=bio,
                attitude=attitude,
                last_message=last_message,
            ).model_dump(),
            "output": output,
        })

        # Call helper library
    create_examples(dict_list, dataset_name=LANGSMITH_DATASET_NAME)

def get_sheet_values(sheet_range: str) -> list[Any]:
    credentials = service_account.Credentials.from_service_account_file(
        SERVICE_ACCOUNT_KEY_PATH,
        scopes=["https://www.googleapis.com/auth/spreadsheets.readonly"],
    )

    service = build("sheets", "v4", credentials=credentials)

    sheet = service.spreadsheets()
    result = sheet.values().get(spreadsheetId=SPREADSHEET_ID, range=sheet_range).execute()
    values = result.get("values")

    return values

if __name__ == "__main__":
    asyncio.run(dataset())

By executing this script (python -m dataset.py), we were able to reflect all inputs and outputs defined in the spreadsheet to the Examples of any Dataset.

Next, we describe the specification of Datasets, LLM models, builtin check functions we want to use, etc. in the YAML file (config.yaml) that defines the config necessary for the Experiment using the Helper.

description: test prompt

prompt:
  name: prompt.py
  type: python
  entry_function: prompt

evaluators_file_path: evaluations.py

providers:
  - id: TURBO

tests:
  type: langsmith_db
  dataset_name: Test Dataset
  experiment_prefix: test_

By the way, if you want to apply custom checks, you can easily run that evaluation by writing evaluation logic in evaluations.py and specifying it in the config.

For example, we were able to implement and use custom check functions for accuracy judgment of positive and negative for prompts that return scores including positive and negative.

def same_sign(run: Run, example: Example) -> dict[str, bool]:
    if run.outputs is None or example.outputs is None or run.outputs.get("output") is None:
        return {"score": False}

    example_output = example.outputs.get("output")
    if example_output is None:
        raise ValueError("example.outputs is None")

    example_score = int(example_output)
    output_score = int(run.outputs["output"].score)
    if example_score == 0 or output_score == 0:
        return {"score": example_score == output_score}

    example_sign = example_score // abs(example_score)
    output_sign = output_score // abs(output_score)
    return {"score": example_sign == output_sign}

This wasn’t in the initial requirements, but we were taught that such things can be used as Experiments functions.

The score results of the implemented check functions can be visually displayed in Charts on LangSmith’s Experiments, which was also very easy to understand as an evaluation indicator.

Finally, we define the implementation of the prompt template to be executed in a separate file (prompt.py), referring to the actual production-side prompt template definition. (It’s also possible to define the prompt text directly on this file)

@traceable
async def prompt(inputs: dict[str, Any]) -> TestPromptTemplate:
    return TestPromptTemplate(
        input=TestPromptInput(
            name=inputs.get("name", ""),
            bio=inputs.get("bio", ""),
            attitude=inputs.get("attitude", ""),
        )
    )

Then, when you execute the alias command that refers to the config, LLM execution is performed for the number of registered datasets, and after execution is complete, all execution results are saved in Experiments on LangSmith.

make evaluate <path/to/config.yml>

Since there are URL links for each of these Experiments results, sharing this link made it easy for tuning personnel and stakeholders to discuss execution results and modification policies 🎉

Operation Phase

As mentioned earlier, the actual prompt tuning task involved multiple prompts, so we divided the work, including the above implementation, among engineer members (basically, one person was in charge of one prompt each).

We summarized in Notion, for each prompt, what kind of logic was used to correct phenomena that engineers could tune and were clearly wrong against the expected value, so that we could track the content of tuning chronologically.

Also, when we got good results to some extent, we shared the Output LangSmith Experiments URL with Biz members for review, continued to make corrections based on their feedback, or completed if there were no problems.

Results and Considerations

Previously, we had an unfavorable situation where engineers and PdMs familiar with LLM development were stuck on prompt tuning for about two weeks, repeatedly reviewing and modifying prompt improvements, or sometimes the tuning itself was neglected.

With the efficiency improvement of prompt tuning, we were able to achieve results where the data creation work in spreadsheets centered on Biz took a few days, the implementation of Datasets and Evaluation Script by engineers took 0.5 to 1 day, and the tuning and review implementation took a few days depending on the complexity of the prompt (actually, we were able to complete the tuning implementation period in about a week).

This achievement is expected to be quite good information as a guideline for doing similar work in the future.

Also, while previously only specific members performed this task, having all engineers in the team share this process had secondary effects such as acquiring context and knowledge of prompt tuning and mutual review.

Until now, we had Biz members actually touch the completed application, perform qualitative quality evaluation, and list modification perspectives. This time, by having them register the expected output values as a dataset in advance as ground truth, we were able to greatly reduce the psychological burden on the tuning side.

We gave Git links to the prompt definition files to Biz members in advance, explained and had them understand the actual prompt text itself, how it works, and in what situations the relevant prompt is executed, before conducting tuning. As it became easy to view Chain trace data from LangSmith’s Experiments results, we were able to receive suggestions from Biz members for modification policies, such as “Wouldn’t it be better to change this notation in this prompt like this?” As engineers, it’s really helpful to receive such detailed perspectives.

Prompt tuning work is directly linked to quality and cannot be completed by engineers alone. We also gained the realization that it’s very important for context to mutually permeate including surrounding members.

Future Challenges

Our efficiency improvements are still quite rough and ad-hoc improvements at present.

We expect various challenges to arise as we continue to develop LLM products in the future, but here are some examples of future challenges and the improvements we are considering.

In particular, when it comes to mass-producing character AIs in the future, the human power cost aspect of Biz members is a very big challenge, so we are considering approaching this next.

• Human power cost of creating feature and memory data for creating character AI
   — Automatically generate feature and memory data from personality definition data using generative AI
• Human power cost of creating ground truth data as expected Output values
   — Infer and automatically generate correct data for new AI based on existing datasets and personality definition data using generative AI
• Quality issues and cost issues of parts where prompt accuracy Evaluation and tuning are done by human power
   — Consider mechanisms for AI to evaluate and tune using LLM as a Judge

Conclusion

I introduced an approach to improve efficiency using Google Spreadsheet and LangSmith’s Helper for challenges that arise in the rather common task of prompt tuning.
I feel we were able to create one case of how to improve efficiency in labor-intensive and time-consuming tasks, and how to apply new technology like generative AI to resolve it.
We hope to continue improving and make LLM product development better.

That’s all, thank you for reading.

More Articles:

Gaudiy AI Lab

About Company:

Gaudiy Inc. | Moving forward in a new era with fans

Part 2: Process— How We Used LangSmith to Streamline Evaluation and Experimentation in LLM Product… was originally published in Gaudiy AI Lab on Medium, where people are continuing the conversation by highlighting and responding to this story.