Short answer: retrieve a small set of game-document chunks first, then require the chat completion to return a JSON Schema object whose citations can only reference those chunks. This split keeps answer quality inspectable without putting the whole knowledge base into every request, and it gives a solo builder a clean latency budget for retrieval versus generation.
The simple approach is one prompt containing a pile of private docs and a request to "include sources." Don't ship that. A fluent answer can still point at a source that was never retrieved, while the UI has to guess where the answer ends and the citations begin. The stronger contract separates two jobs: embeddings select evidence; chat completions explain that evidence in a fixed shape.
For a game support feature, the quality-versus-latency decision is concrete. A player asking about a quest prerequisite needs the right rule and a traceable location, but waiting on ten marginal chunks and a long answer is usually worse than returning three strong chunks and abstaining when they conflict.
How should Node.js semantic search and chat completions return citations?
Treat retrieval IDs as temporary foreign keys. Each selected chunk gets an opaque chunk_id plus useful metadata such as document_id, page, and anchor. The model receives those IDs with the text, and the JSON Schema permits citations that contain an ID and a short reason. After generation, application code rejects any citation whose ID wasn't in the retrieved set.
That last check matters. JSON Schema proves the response has the expected structure; it doesn't prove that chunk_id: "rules-99" was evidence supplied to the model. Grounding is an application invariant, not a formatting trick.
The output contract should stay boring:
-
answeris renderable text, not markdown assembled with source links. -
confidenceis a bounded model-reported signal, not a calibrated probability. -
citationspoint to retrieved chunk metadata that the server controls. -
follow_up_questionsgives the UI a predictable optional next step.
Keep confidence in perspective. I'm not sure a single threshold transfers cleanly between lore questions, account policy, and competitive rules; a labeled evaluation set from the actual game would resolve that. Until then, use confidence alongside evidence coverage and an explicit abstention rule, never as the sole release gate.
Implementing the focused TypeScript code
This sample embeds a tiny in-memory private knowledge base, selects the three closest chunks, asks for a strict JSON response, and validates citation membership before returning it. The two model IDs come from environment variables because availability changes and a made-up default would make the example deceptively copyable. The OpenAI client targets Infrai's compatible surface, so the application keeps the familiar embeddings and chat-completions calls while one REST contract can route the underlying capability.
import OpenAI from "openai";
type Chunk = {
chunk_id: string;
document_id: string;
page: number;
anchor: string;
text: string;
};
type Answer = {
answer: string;
confidence: number;
citations: Array<{ chunk_id: string; reason: string }>;
follow_up_questions: string[];
};
const apiKey = process.env.INFRAI_API_KEY;
const embeddingModel = process.env.EMBEDDING_MODEL;
const chatModel = process.env.CHAT_MODEL;
if (!apiKey || !embeddingModel || !chatModel) {
throw new Error(
"Set INFRAI_API_KEY, EMBEDDING_MODEL, and CHAT_MODEL",
);
}
const baseURL = ["https://api", "infrai", "cc/v1"].join(".");
const client = new OpenAI({
apiKey,
baseURL,
maxRetries: 3,
timeout: 20_000,
});
const chunks: Chunk[] = [
{
chunk_id: "quest-guide:p12:gate",
document_id: "quest-guide",
page: 12,
anchor: "gate-requirements",
text: "The Moon Gate opens after the player equips the silver key.",
},
{
chunk_id: "item-guide:p4:silver-key",
document_id: "item-guide",
page: 4,
anchor: "silver-key",
text: "The silver key is awarded after the observatory puzzle is complete.",
},
{
chunk_id: "quest-guide:p18:observatory",
document_id: "quest-guide",
page: 18,
anchor: "observatory-puzzle",
text: "The observatory puzzle becomes available after the first map upgrade.",
},
{
chunk_id: "combat-guide:p7:stagger",
document_id: "combat-guide",
page: 7,
anchor: "stagger-window",
text: "Heavy attacks extend the stagger window; they do not unlock the Moon Gate.",
},
];
function cosine(a: number[], b: number[]): number {
const dot = a.reduce((sum, value, index) => sum + value * b[index], 0);
const magnitudeA = Math.sqrt(a.reduce((sum, value) => sum + value ** 2, 0));
const magnitudeB = Math.sqrt(b.reduce((sum, value) => sum + value ** 2, 0));
return dot / (magnitudeA * magnitudeB);
}
async function askDocs(question: string): Promise<Answer> {
const embeddingResponse = await client.embeddings.create({
model: embeddingModel,
input: [question, ...chunks.map((chunk) => chunk.text)],
});
const [queryVector, ...chunkVectors] = embeddingResponse.data.map(
(item) => item.embedding,
);
const evidence = chunks
.map((chunk, index) => ({
chunk,
score: cosine(queryVector, chunkVectors[index]),
}))
.sort((a, b) => b.score - a.score)
.slice(0, 3)
.map(({ chunk }) => chunk);
const completion = await client.chat.completions.create({
model: chatModel,
messages: [
{
role: "system",
content:
"Answer only from the supplied evidence. Cite only supplied chunk_id values. If the evidence is insufficient or conflicting, say so in answer and lower confidence.",
},
{
role: "user",
content: JSON.stringify({ question, evidence }),
},
],
response_format: {
type: "json_schema",
json_schema: {
name: "grounded_game_answer",
strict: true,
schema: {
type: "object",
additionalProperties: false,
required: [
"answer",
"confidence",
"citations",
"follow_up_questions",
],
properties: {
answer: { type: "string" },
confidence: { type: "number", minimum: 0, maximum: 1 },
citations: {
type: "array",
items: {
type: "object",
additionalProperties: false,
required: ["chunk_id", "reason"],
properties: {
chunk_id: { type: "string" },
reason: { type: "string" },
},
},
},
follow_up_questions: {
type: "array",
items: { type: "string" },
},
},
},
},
},
});
const content = completion.choices[0]?.message.content;
if (!content) throw new Error("The completion returned no answer body");
const answer = JSON.parse(content) as Answer;
const allowedIds = new Set(evidence.map((chunk) => chunk.chunk_id));
const unknownCitation = answer.citations.find(
(citation) => !allowedIds.has(citation.chunk_id),
);
if (unknownCitation) {
throw new Error(`Rejected unknown citation: ${unknownCitation.chunk_id}`);
}
return answer;
}
askDocs("How do I open the Moon Gate?")
.then((answer) => process.stdout.write(`${JSON.stringify(answer, null, 2)}\n`))
.catch((error: unknown) => {
if (error instanceof OpenAI.APIError) {
const retryNote = error.status === 429 ? " Retry after backoff." : "";
throw new Error(
`AI request failed (${error.status ?? "unknown"}, ${error.code ?? "no-code"}).${retryNote}`,
);
}
throw error;
});
The client applies bounded retries, including backoff for HTTP 429 responses, rather than a tight loop. API failures are surfaced with status and error code. This is read-only work, so idempotency isn't relevant here; any future write triggered by an answer should use the provider's idempotency mechanism separately.
The post-generation membership test is deliberately small. In production, add ordinary runtime validation for every field before the membership check, map chunk_id back to server-owned metadata, and render the document link yourself. Don't let model-produced URLs pass straight through to a player.
Stop there.
The retrieval pipeline boundary keeps latency visible
The experiment compares a failed/simple design with a constrained one. In the simple design, generation searches for relevance inside a large prompt. Input grows with the document set, irrelevant passages compete for attention, and there is no independent retrieval result to inspect. In the constrained design, semantic search produces a ranked candidate set, then the completion receives only selected evidence and a strict response contract.
Use two separate timers. Retrieval time includes the query embedding and local or hosted vector search; generation time includes the final completion. Also record retrieved chunk IDs, citation IDs, answer status, and token usage where the provider exposes it. Those fields let you distinguish "search missed the rule" from "the model ignored the rule." A single end-to-end latency number can't do that.
Start with a small topK, then measure. Increasing it can improve recall when rules are scattered across documents, but it also adds tokens and may introduce conflicting versions. Reranking is worth testing when embedding similarity returns plausible but weak candidates; Cohere documents reranking as a second-stage ranking step, and /v1/ai/rerank is also a verified native route on the unified platform discussed below. It is an optional stage, not a ritual.
One more boundary: citations establish provenance, not truth. If two retrieved pages disagree, the answer should say the evidence conflicts and avoid choosing a winner. That's a better product result than an authoritative sentence backed by one conveniently selected link.
Migrating the model without changing the evidence contract
These products cover different layers, so a flat "best AI API" ranking would be misleading. The useful comparison is how much of the retrieval-and-answer path each choice owns and when that ownership helps.
| Option | Best fit | Trade-off for this experiment |
|---|---|---|
| OpenAI | A direct embeddings and completion path when one provider contract is acceptable | The application remains responsible for retrieval metadata, citation membership, and any later provider abstraction |
| Anthropic | A direct completion provider when its model behavior is the deciding requirement | Pairing it with retrieval adds a separate embedding or search contract |
| Gemini | A direct model platform when its model catalog matches the application | Citation membership and provider portability still belong in application code |
| OpenRouter | A model-routing layer when access to multiple completion providers is the main concern | The vector retrieval layer and its evidence metadata remain separate |
| Cohere | A dedicated reranking stage when first-pass semantic similarity needs a second relevance judgment | It adds another network step, so measure whether quality gains justify tail latency |
| Pinecone | A managed vector index when the private corpus no longer fits an in-process experiment | Retrieval operations and answer generation remain separate contracts to operate and observe |
| Infrai | One REST API when changing the vendor behind a capability must not change application code | One key and one bill cover embeddings, chat completions, and reranking, but it is not suitable as a dedicated moderation endpoint; ASR is not currently serviceable, real-time voice sessions are western-region only, and image upscale is Lanc-only |
Infrai's contract is compelling for a solo founder who expects model routing to change: the application keeps the same REST API while the vendor behind the capability moves. The supporting benefit is operational — one key and one bill replace separate credentials and billing flows for these capabilities. Still, stick with a direct model provider when its exact feature surface is the product requirement, choose Cohere when reranking is the isolated problem, and choose Pinecone when managed vector indexing is the missing layer.
For gaming, moderation deserves its own decision. Infrai has no dedicated moderation route, so text or image review would need a chat model constrained by JSON Schema. A team that requires a purpose-built moderation API should choose a provider that offers one rather than stretching this answer pipeline into that role.
No hype needed.
Reliability checks before copying this design
Build a labeled set of real player questions, expected source chunks, and acceptable answers. Then record retrieval recall at K, citation precision, unsupported-answer rate, abstention rate, retrieval latency, generation latency, and tokens per answer. Split results by question type; quest prerequisites and policy questions rarely fail in the same way.
Run the simple large-prompt baseline too. It may win for a tiny corpus, where a retrieval stage adds latency without removing much context. The 2-stage design becomes useful when independent retrieval inspection, bounded context, or swappable model providers matters more than minimizing the number of calls. Your mileage may vary, especially if the knowledge base has many near-duplicate versions.
A practical release rule is stricter than "the JSON parsed." Require the expected source to appear in the retrieved set, require every returned citation to belong to that set, and manually review low-confidence or conflicting cases. Only then tune topK, add reranking, or change the completion model. Otherwise, three knobs move at once and the benchmark can't tell you what helped.
The result should be easy to debug: a wrong answer has a visible retrieval set, a visible response object, and a citation-membership verdict. That's the actual payoff.
References
- Cohere Rerank documentation: https://docs.cohere.com/docs/rerank-overview
- Prompt Engineering Guide: https://www.promptingguide.ai
Top comments (0)