Bringing trustworthy, AI-powered answers into your Salesforce org is now much simpler — thanks to tools like Apex, prompt templates, and Agentforce. In this post, we’ll explore how you can create an AI agent that uses real web search results to answer questions, with citations included right in the response. Whether you’re building something for your team or for your customers, this approach helps make your AI more reliable by showing where its answers come from.
Setting up the prompt template: The AI brain behind the response
Note: Before getting started, you’ll first have to set up a web retriever in Data Cloud. You can set this up from Data Cloud > Einstein Studio > Retrievers. See more in Salesforce Help.
Next, prior writing any Apex code, you’ll need to create the prompt template that defines how the AI should answer questions. This template acts as the brain of your system — it handles the actual retrieval, reasoning, and formatting of the response.
When setting it up in Prompt Builder, be sure to take note of two important details:
- The API name of the prompt template
- The name of the input variable that sends the user’s question
These two values will be referenced later in the Apex class, so they need to match exactly. For example, you can name the template Knowledge_Search and the input variable Question (we’ll use these in the Apex code shown in the next section of this post).
Here’s an example of what the prompt content might look like inside the template:
Use the following web results to answer the question. Provide citations for any facts you use. If more than 1 results, select the one(s) with more credibility. Web Results: {!$EinsteinSearch:sfdc_ai__WebRetrievalAction.results} Question: {!$Input:Question}Once your template is set up and saved, don’t forget to activate it. That final step makes it available for Apex and Agentforce to use.
Connecting Apex to the prompt template
With the prompt template in place, it’s time to build the Apex class that connects everything together. This Apex class will be used as an agent action inside Agentforce. It sends the user’s question to the prompt template via the Connect API, collects the AI’s response (including citations), and transforms the results into a format that your agent needs to show citations.
// Invokes Einstein Prompt Template via Connect API, gets response/citations, transforms citations. public class KnowledgeBaseApex { // Calls 'Knowledge_Search' Prompt Template via Connect API, gets raw response/citations. public static ConnectApi.EinsteinPromptTemplateGenerationsRepresentation executePromptTemplate(String inputQuestion) { ConnectApi.WrappedValue inputText = new ConnectApi.WrappedValue(); inputText.value = inputQuestion; Map<String, ConnectApi.WrappedValue> inputParams = new Map<String, ConnectApi.WrappedValue>(); // NOTE: 'Input:Question' must match the Prompt Template input variable API name. inputParams.put('Input:Question', inputText); ConnectApi.EinsteinPromptTemplateGenerationsInput executeTemplateInput = new ConnectApi.EinsteinPromptTemplateGenerationsInput(); executeTemplateInput.additionalConfig = new ConnectApi.EinsteinLlmAdditionalConfigInput(); executeTemplateInput.additionalConfig.applicationName = 'PromptBuilderPreview'; executeTemplateInput.isPreview = true; executeTemplateInput.citationMode = 'post_generation'; executeTemplateInput.inputParams = inputParams; ConnectApi.EinsteinPromptTemplateGenerationsRepresentation generationsOutput = ConnectApi.EinsteinLLM.generateMessagesForPromptTemplate( 'Knowledge_Search', // Developer Name of the Prompt Template executeTemplateInput ); return generationsOutput; } // Helper methods to transform Connect API citations to AI Copilot format. // Transforms a single Connect API source reference to AiCopilot format. public static AiCopilot.GenAiSourceReference transformSourceReference(ConnectApi.EinsteinLlmGenAiSourceReference sourceRef) { List contents = new List(); for (ConnectApi.EinsteinLlmGenerationGenAiSourceContentInfo contentInfo : sourceRef.contents) { contents.add(new AiCopilot.GenAiSourceContentInfo(contentInfo.fieldName, contentInfo.objectName, contentInfo.content)); } List metadata = new List(); for (ConnectApi.EinsteinLlmGenerationGenAiSourceReferenceInfo referenceInfo : sourceRef.metadata) { metadata.add(new AiCopilot.GenAiSourceReferenceInfo(referenceInfo.link, referenceInfo.sourceObjectRecordId, referenceInfo.sourceObjectApiName)); } return new AiCopilot.GenAiSourceReference(null, contents, metadata); } // Transforms the overall Connect API citation block to AiCopilot format. public static AiCopilot.GenAiCitationInput transformCitations(ConnectApi.EinsteinLlmGenerationCitationOutput citations) { List sourceRefs = new List(); if (citations != null && citations.sourceReferences != null) { for (ConnectApi.EinsteinLlmGenAiSourceReference source : citations.sourceReferences) { sourceRefs.add(transformSourceReference(source)); } } return new AiCopilot.GenAiCitationInput(null, sourceRefs); } // Defines input (Request) and output (Response) structures for the Invocable Method. public class Request { @InvocableVariable(required=true description='Complete user question') public String Question; } public class Response { @InvocableVariable(required=true description='Generated text response from the LLM') public String Data; @InvocableVariable(required=true description='Structured citation sources') public AiCopilot.GenAiCitationInput sources; } // @InvocableMethod for Flows: gets question, calls execute, transforms citations, returns results. @InvocableMethod(label='Apex executePrompt from KnowledgeBaseApex' description='Invokes the Knowledge Search prompt template and returns the response with citations.') public static List executePrompt(List requests) { String question = requests[0].Question; ConnectApi.EinsteinPromptTemplateGenerationsRepresentation output = executePromptTemplate(question); List responses = new List(); Response response = new Response(); response.Data = output.prompt; response.sources = transformCitations(output.citations); responses.add(response); return responses; } }What each part of the Apex code does
executePrompt: This is the main method that Agentforce will call. It takes the input question, sends it to the prompt template, transforms the citations, and returns everything — ready to be shown to the user.executePromptTemplate: This method takes the user’s question and sends it to a prompt template calledKnowledge_Searchusing the Connect API. It formats the input correctly, tells the AI to include citations, and returns everything the AI sends back, including the answer and the raw citation info. Note: The input variable name'Input:Question'must exactly match what’s defined in the prompt template.transformSourceReference: This method takes one source (a single citation) from the AI response and turns it into the format that Agentforce expects. It pulls out both the content (what the source says) and the metadata (like where it came from).transformCitations: This method looks at the full set of citations from the AI and transforms each one using the helper method above. It returns everything in a nice, structured format that can be passed back into Agentforce.Requestclass: This is the input model. It defines what the agent or flow sends into the Apex method; in this case, just the user’s full question as a single string.Responseclass: This is the output model. It defines what comes back from the Apex method: the AI’s answer and the list of sources (citations) in a structured format.
Connecting the pieces inside your agent
Once the Apex code and prompt template are ready, the last step is to wire everything together in Agentforce. This happens inside the topic, where you define the kind of questions your agent will handle.
In this case, you’ll create a topic focused on web information retrieval. The agent uses this context to search external sources and bring back credible, cited answers. Within that topic, you’ll link your Apex class as an Apex action — this is what triggers the AI workflow when a user asks a question.
The prompt template will continue to handle how the response is formatted, so there’s no need to surface raw data or sources separately in the agent UI. It all comes back clean and ready to present.
Conclusion
By bringing together Apex, prompt templates, and Agentforce, you’ve built a Salesforce-native agent that delivers reliable, cited answers to complex questions. This approach is especially valuable for knowledge-heavy scenarios, such as internal support, finding documentation, or navigating regulations, where accuracy and transparency really matter.
To take it further, you might connect this agent to customer-facing channels, add role-based access controls, or combine web search with document-based retrieval for even richer responses.
Resources
- Documentation: Get Started with Citations in Apex
- Agentforce Decoded Video: Build a Cited Web Search Agent with Agentforce
About the author
Alex Martinez was part of the MuleSoft Community before joining MuleSoft as a Developer Advocate. She founded ProstDev to help other professionals learn more about content creation. In her free time, you’ll find Alex playing Nintendo or PlayStation games and writing reviews about them. Follow Alex on LinkedIn or in the Trailblazer Community.
