Skip to content

RAG AgentS: Simple In-Context RAG

One significant task that is needed in the total implementation is to link a natural language description to a specific task to perform in the implementation. ChemConnect uses ontologies to describe all the data objects and the tasks used in ChemConnect (see publications 2019 and 2021). Associated with each object, whether it be a data object or a task, there is a description (in the dcterms:comment in the annotations of the object) of the object. 

If the user specifies a task using natural language, the specific task is to find the description of the object that is the most similar. This the basis of a simple RAG implementation. Given the pair:

  • Keyword: This is the name of the ontology object. It could be a data type or a task or any object defined in the ChemConnect ontology.
  • Description: This is a natural language description of the object. This is found in the annotations of the object under dcterms:comment.

The goal is to match the natural language prompt with the specific keyword representing the specific ontology object. This ontology object then guides how to proceed. 

The goal of this prototype is to interpret the prompt and use Retrieval Augmented Generation (RAG) methodology determine which keyword it is describing. In this prototype, in-context RAG is being used, meaning that the keyword and its description is directly in the prompt. 

The base of the prototype will start JSON object with the keyword and the description. The source of this set of pairs can be from directly defined JAVA JsonObjects to generated from the ChemConnect ontology. The first prototype will be a small set of specifically defined objects.

Simple RAG Prototype

The simple RAG prototype interprets the prompt and responds with the closest identifier (and its corresponding description). The RAG input is simply a ‘ID’ and a ‘Description’. For example:

  • ID: InitialReadInLocalStorageSystem,
  • Description: Read in text file from the local file system of the client to the staging area.

So if the description is interpreted to be close to that of the RAG element, then the ID is returned.  A confidence is given of the match also. In the prototype this is the response in the chat window. The figure gives several examples.

Simple RAG chat window

SimpleRagElement

The RAG information is stored in text form and given to the agent in the simpleDictionaryString. The section ‘The RAG Data‘ below shows the data and how it is defined.

This agent returns a JSON object with three elements:

  • formalName: This is the formal name of the concept
  • description: This is the description within the RAG that was matched with the prompt input.
  • confidence: This the confidence level of the match.
SimpleRagAgents.java
package esblurock.info.JavaAIAgent.simplerag;
import java.util.List;
import java.util.Map;
import java.util.Vector;
import java.util.stream.Collectors;
import com.google.adk.agents.LlmAgent;
import com.google.genai.types.Schema;
import com.google.gson.JsonObject;
import esblurock.info.JavaAIAgent.Constants;
import esblurock.info.JavaAIAgent.simplerag.dictionary.DictionaryElementSet;
import esblurock.info.JavaAIAgent.simplerag.dictionary.SimpleDictionarySet;

public class SimpleRagAgents {
	
	public static String simpleDictionaryString = SimpleDictionarySet.createSimple().toRAGString();
	public static LlmAgent EXTRACTOR_AGENT = LlmAgent.builder()
	    .name("ParameterExtractor")
	    .model(Constants.MODEL)
	    .instruction("You are a specialized entity extractor. Your goal is to map user text " +
	                 "to a 'formal name' from the provided dictionary below.\n\n" +
	                 "DICTIONARY:\n" + simpleDictionaryString + "\n\n" +
	                 "RULES:\n" +
	                 "1. Analyze the user's input text in the prompt field of the input JSON object.\n" +
	                 "2. Identify which parameter description best matches the user's intent.\n" +
	                 "3. Return the formal name and the description of the matching element as the result.\n" +
	                 "4. If no clear match exists, return 'UNKNOWN'.")
	    // Use an output schema to ensure you get a clean JSON object back
	    .outputSchema(Schema.builder()
	        .type("OBJECT")
	        .properties(Map.of(
		            "formalName", Schema.builder().type("STRING").description("The matched identifier").build(),
		            "description", Schema.builder().type("STRING").description("The description of the identifier").build(),
	            "confidence", Schema.builder().type("STRING").description("HIGH, MEDIUM, or LOW").build()
	        ))
	        .required(List.of("formalName"))
	        .build())
	    .build();
}

The RAG Data

The origin of the RAG data in this prototype is a JSON object. The information in the JSON object  will be converted to text to be inserted in the prompt. This is in-context RAG. 

The use of a JSON file as the origin of the information is twofold, one the individual elements of data can be accessed and queried. The other is that this leaves open the origin of the data. Within the ChemConnect implementation the origin will ultimately come from the ontology. Specifically, the ontology class name as the RAG label and its description as the RAG natural language description.

The JSON schema is held in two classes: 

  • DictionaryElement: This is the individual element of the RAG with the ID and the description. The setters, getters and constructors are defined.
  • DictionaryElementSet: This holds the set of DictionaryElements. This has the constructors, addElement and toRAGString routines.

The JSON file is converted to a string through the toRAGString method in the DictionaryElementSet. The agent prompt needs the string to like the following example (this is RAG set used in this example):

- ID: TransactionInterpretDisassociationEnergy, Description: Interpret a parsed file with disassociation text blocks to dataset objects. For each definition the structure in short text format, a structure name and the disassociate energy value.
- ID: TransactionInterpretMetaAtom, Description: Interpret a parsed file of single line definiitions of Meta Atom dataset objects. Each line has a Meta Atom name and a short text line structure, structure name, meta atom name and meta atom type.
- ID: TransactionInterpretSymmetryInformation, Description: Interpret a parsed file of symmetry definitions in JThermodynamics XML structure to dataset objects. The symmetry object has a molfile structure and assignments to symmetry.
- ID: TransactionInterpretThermodynamicBlockLines, Description: From a single (3 or 4 line) block of the thermodynamic specification, the thermodynamics and structure are isolated out and interpreted from a parsed file. The format originates from the TherGas format for defining thermodynamics.
- ID: TransactionInterpretVibrationalMode, Description: Interpret a parsed file of one line vibrational definitions to dataset objects. Each line has the vibrational structure in short text form, vibrational name, structure name, vibrational value and symmetry factor.
- ID: PartiionSetWithinRepositoryFile, Description: The repository file represents the set of inputs for a set of catalog dataset objects. This parses the file into individual partitions.
- ID: InitialReadInLocalStorageSystem, Description: Read in text file from the local file system of the client to the staging area.

DictionaryElement

This class defines the schema for RAG information, holding the single keyword of the object (‘term’) and the description of the object (‘description’)

@JsonProperty(“term”)

private String term;

@JsonProperty(“descriptiion”)

private String description;

 

DictionaryElement.java
package esblurock.info.JavaAIAgent.simplerag.dictionary;
import com.fasterxml.jackson.annotation.JsonProperty;
public class DictionaryElement {
	@JsonProperty("term")
    private String term;
	@JsonProperty("descriptiion")
    private String description;
	
	public DictionaryElement(String term, String description) {
		this.term = term;
		this.description = description;
	}
	public String getTerm() {
		return term;
	}
	public void setTerm(String term) {
		this.term = term;
	}
	public String getDescription() {
		return description;
	}
	public void setDescription(String description) {
		this.description = description;
	}
}

DictionaryElementSet

This is the JAVA class representation of the RAG information. Basically it is a list of DictionaryElement objects. In a ‘real’ system, this JAVA class could be filled with the RAG information, regardless of source. The method within this class,

toRAGSTring(),

gives a string that can be inserted into the agent managing the particular RAG. The constructor

DictionaryElementSet(String dictionaryName, String description)

sets up set with the name and description. RAG elements are added to the dictionary using the method:

addElement(DictionaryElement element)

DictionaryElement is used as the input to allow changes to the fundamental data that is found

There are three elements in the DictionaryElementSet:

  • dictionaryName: This give a name to the RAG information set.
  • dictionaryDescription: This is a short description of the type of information that is contained within this set
  • elements: This is the actual list of DictionaryElement objects.

getters and setters methods are available for each of these.

 

filename.java
package esblurock.info.JavaAIAgent.simplerag.dictionary;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import com.fasterxml.jackson.annotation.JsonProperty;
public class DictionaryElementSet {
	@JsonProperty("dictionaryName")
    private String dictionaryName;
	@JsonProperty("dictionaryDescriptiion")
    private String description;
	@JsonProperty("dictionaryDescriptiion")
	private List<DictionaryElement> elements;
	
	public DictionaryElementSet(String dictionaryName, String description, List<DictionaryElement> elements) {
		this.dictionaryName = dictionaryName;
		this.description = description;
		this.elements = elements;
	}
	
	public DictionaryElementSet(String dictionaryName, String description) {
		this.dictionaryName = dictionaryName;
		this.description = description;
		this.elements = new ArrayList<DictionaryElement>();
	}
	public String getDictionaryName() {
		return dictionaryName;
	}
	public void setDictionaryName(String dictionaryName) {
		this.dictionaryName = dictionaryName;
	}
		
	public String getDescription() {
		return description;
	}
	public void setDescription(String description) {
		this.description = description;
	}
	public List<DictionaryElement> getElements() {
		return elements;
	}
	public void setElements(List<DictionaryElement> elements) {
		this.elements = elements;
	}
	public void addElement(DictionaryElement element) {
		this.elements.add(element);
	}
	public String toRAGString() {
		String dictionaryString = this.elements.stream()
				.map(obj -> String.format("- ID: %s, Description: %s", obj.getTerm(), obj.getDescription()))
				.collect(Collectors.joining("\n"));
		System.out.println("Dictionary String:\n" + dictionaryString);
		return dictionaryString;
	}
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append("Dictionary Name: ").append(dictionaryName).append("\n");
		sb.append("Description: ").append(description).append("\n");
		sb.append("Elements:\n");
		for (DictionaryElement element : elements) {
			sb.append(" - ID: ").append(element.getTerm()).append(", Description: ").append(element.getDescription())
					.append("\n");
		}
		return sb.toString();
	}
}

SimpleDictionarySet

This class sets up the RAG elements with the createSimple() method. This manually sets up a RAG dictionary. This mimics actual ontology information. 

SimpleDictionarySet.java
package esblurock.info.JavaAIAgent.simplerag.dictionary;
import com.google.gson.JsonObject;
public class SimpleDictionarySet {
	
	public static DictionaryElementSet createSimple() {
		String dictionaryNameString = "";
		String descriptionString = "";
		DictionaryElementSet simpleDictionary = new DictionaryElementSet(dictionaryNameString,descriptionString);
		
		DictionaryElement e1 = new DictionaryElement(
				"TransactionInterpretDisassociationEnergy", 
				"Interpret a parsed file with disassociation text blocks to dataset objects. For each definition the structure in short text format, a structure name and the disassociate energy value.");
		DictionaryElement e2 = new DictionaryElement(
				"TransactionInterpretMetaAtom", 
				"Interpret a parsed file of single line definiitions of Meta Atom dataset objects. Each line has a Meta Atom name and a short text line structure, structure name, meta atom name and meta atom type.");
		DictionaryElement e3 = new DictionaryElement(
				"TransactionInterpretSymmetryInformation", 
				"Interpret a parsed file of symmetry definitions in JThermodynamics XML structure to dataset objects. The symmetry object has a molfile structure and assignments to symmetry.");
		DictionaryElement e4 = new DictionaryElement(
				"TransactionInterpretThermodynamicBlockLines", 
				"From a single (3 or 4 line) block of the thermodynamic specification, the thermodynamics and structure are isolated out and interpreted from a parsed file. The format originates from the TherGas format for defining thermodynamics.");
		DictionaryElement e5 = new DictionaryElement(
				"TransactionInterpretVibrationalMode", 
				"Interpret a parsed file of one line vibrational definitions to dataset objects.  Each line has the vibrational structure in short text form, vibrational name, structure name, vibrational value and symmetry factor.");
		DictionaryElement e6 = new DictionaryElement(
				"PartiionSetWithinRepositoryFile", 
				"The repository file represents the set of inputs for a set of catalog dataset objects. This parses the file into individual partitions.");
		DictionaryElement e7 = new DictionaryElement(
				"InitialReadInLocalStorageSystem", 
				"Read in text file from the local file system of the client to the staging area.");
		
		simpleDictionary.addElement(e1);
		simpleDictionary.addElement(e2);
		simpleDictionary.addElement(e3);
		simpleDictionary.addElement(e4);
		simpleDictionary.addElement(e5);
		simpleDictionary.addElement(e6);
		simpleDictionary.addElement(e7);
		
		return simpleDictionary;
	}
}

SimpleRAGServlet

The SimpleRAGServlet JAVA class uses the generic servlet defined earlier  to set up the agent with the application name (Constants.SIMPLERAGAPPNAME) and the link to the RAG agent (SimpleRagAgents.EXTRACTOR_AGENT).

filename.java
package esblurock.info.JavaAIAgent.simplerag;

import esblurock.info.JavaAIAgent.Constants;
import esblurock.info.JavaAIAgent.generic.SimpleGenericServlet;
import jakarta.servlet.annotation.WebServlet;
@WebServlet("/api/agent/simplerag")
public class SimpleRagServlet extends SimpleGenericServlet {
	private static final long serialVersionUID = 4749598458122810112L;
	public SimpleRagServlet() {
		super(Constants.SIMPLERAGAPPNAME, SimpleRagAgents.EXTRACTOR_AGENT);
	}
}

Client:  SimpleRAGChat

The SimpleRAGChat component sets up a standard chat window (as in previous examples, such as the simple chat). 

When the prompt is entered, the 

getPrompt()

method is called. It sets up the body to be sent to the backend agent:

  • status: The initial status. This is not used in the current agent, but kept for more general uses.
  • prompt: This is the prompt text
  • data: This is empty, also not used in the current agent, but kept for more general uses.
  • sessionId: The session. It is a generated random number, meaning that each prompt is a separate session.
  • userId: a constant representing the current user.

The body is then sent to the 

onSend(body:any)

method. This sends the body to the backend (through the agentService.sendMessage method). The response is interpreted where the three JSON components are extracted:

  • formalName: This is the formal name of the concept
  • description: This is the description within the RAG that was matched with the prompt input.
  • confidence: This the confidence level of the match.

These are combined in a string

const text = `The match is to ${formalName}: ${description} (confidence ${confidence})`;

and send to the chat window:

this.agentService.addMessage({ role: ‘assistant’, text: text });

📘
simple-ragchat.ts
import { Component } from '@angular/core';
import { AgentGeneric } from '../../services/agent-generic';
import { CommonModule } from '@angular/common';
import { FormsModule } from '@angular/forms';
import { MatCardModule } from '@angular/material/card';
import { MatFormFieldModule } from '@angular/material/form-field';
import { MatInputModule } from '@angular/material/input';
import { MatButtonModule } from '@angular/material/button';
import { MatIconModule } from '@angular/material/icon';
@Component({
  selector: 'app-simple-ragchat',
  imports: [
	CommonModule,
	FormsModule,
	MatCardModule,
	MatFormFieldModule,
	MatInputModule,
	MatButtonModule,
	MatIconModule
  ],
  templateUrl: './simple-ragchat.html',
  styleUrl: './simple-ragchat.scss',
})
export class SimpleRAGChat {
	// The prompt field bound to the input box
	userInput = '';
	
	constructor(public agentService: AgentGeneric) {
		// Initialize user and session by the agent
		this.agentService.setUserAndSession();
	}
	
	onSend(body: any) {
		// Send the body to the backend to get a response and add it to the chat
		this.agentService.sendMessage(body,'/api/agent/simplerag').subscribe({
			next: (res) => {
				try {
					const data = JSON.parse(res);
					const formalName = data.formalName;
					const confidence = data.confidence;
					const description  = data.description;
					const text = `The match is to ${formalName}: ${description}  (confidence ${confidence})`;
					this.agentService.addMessage({ role: 'assistant', text: text });
					} catch (e) {
						console.error('Error parsing response', e);
					}
					},
			error: (err) => console.error('Communication failed', err)
					});;
}
getPrompt() {
	// userInput comes from the input field bound with [(ngModel)]
	this.agentService.addMessage({ role: 'user', text: this.userInput });
	const prompt = this.userInput;
	if (!prompt.trim()) return;
	this.userInput = ''; // Reset the input field
	// Add the user's message to the UI immediately through the service
	const data = {};
	const body = {
		statius: 'initial',
		prompt: prompt,
		data: data,
		sessionId: this.agentService.getSessionId(),
		userId: this.agentService.getUserId()
	};
	this.onSend(body);
}
}
🌐
simple-ragchat.html
<mat-card class="chat-container">
	<mat-card-header>
		<mat-card-title>Simple LLM Physics Agent for RAG test</mat-card-title>
		<span class="spacer"></span>
		<button mat-icon-button (click)="agentService.clearChat()">
			<mat-icon>delete_sweep</mat-icon>
		</button>
	</mat-card-header>
	<mat-card-content #scrollContainer class="message-area">
		<div *ngFor="let msg of agentService.messages()">
			<div [ngClass]="msg.role + '-bubble'">
				{{ msg.text }}
			</div>
		</div>
	</mat-card-content>
	<mat-card-actions class="input-area">
		<mat-form-field appearance="outline" class="full-width">
			<mat-label>Ask a question...</mat-label>
			<input matInput [(ngModel)]="userInput" (keyup.enter)="getPrompt()">
			<button mat-icon-button matSuffix (click)="getPrompt()">
				<mat-icon>send</mat-icon>
			</button>
		</mat-form-field>
	</mat-card-actions>
</mat-card>
🎨
filename.scss
.chat-container {
  max-width: 800px;
  margin: 20px auto;
  height: 80vh; /* Give the card a fixed height */
  display: flex;
  flex-direction: column;
}
.spacer { flex: 1 1 auto; }
.message-area {
   height: 400px;
   overflow-y: auto;
   display: flex;
   flex-direction: column;
   padding: 16px;
 }
 .input-area {
    padding: 0 16px 16px 16px;
    display: block; /* Ensures the form-field takes full width */
  }
  .full-width {
    width: 100%;
  }
  .user-bubble {
    align-self: flex-end;
    background-color: #e3f2fd;
    padding: 8px 12px;
    border-radius: 15px 15px 0 15px;
    margin-bottom: 8px;
    max-width: 80%;
  }
  /* User on the right, blueish */
  .assistant-bubble {
    align-self: flex-start;
    background-color: #f5f5f5;
    padding: 8px 12px;
    border-radius: 15px 15px 15px 0;
    margin-bottom: 8px;
    max-width: 80%;
  }

Leave a Reply

Your email address will not be published. Required fields are marked *

en_USEnglish