Improve Embedding Extraction and Add xAI Support #512
Conversation
- Extract text content from multimodal messages for higher quality embeddings. - Update `getEmbeddingModel` to support xAI embeddings while defaulting to OpenAI. - Added documentation for `XAI_EMBEDDING_MODEL` in environment example files. - Ensured robust fallback logic for embedding models. Co-authored-by: ngoiyaeric <115367894+ngoiyaeric@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
|
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Deploying with
|
| Status | Name | Latest Commit | Updated (UTC) |
|---|---|---|---|
| ❌ Deployment failed View logs |
qcx | 7adc7fe | Feb 10 2026, 08:52 AM |
![charliecreates[bot]](https://avatars.githubusercontent.com/in/1140780?s=60&v=4){kind=small}
There was a problem hiding this comment.
The main issues are (1) the new embedding text extraction uses any and silently discards unexpected content shapes, which can hide bugs and reduce retrieval quality, and (2) embedding assignment currently does an O(n²) findIndex loop that can degrade performance for long chats. Additionally, getEmbeddingModel() emits a misleading warning message relative to its actual xAI requirements, making configuration/debugging harder. These are all straightforward to fix without changing the feature behavior.
Additional notes (2)
-
Maintainability |
lib/actions/chat.ts:20-20
CoreMessageis imported but never used in the shown diff. Even if ESLint passes, this is a smell: either the import is accidental or the function should be typed using it (which would be better thanany). -
Maintainability |
.env.local.example:25-29
.env.local.exampleaddsXAI_EMBEDDING_MODELbut notXAI_API_KEY, while.env.exampleincludes both. That inconsistency will cause local setup confusion: users may copy.env.local.exampleexpecting xAI embeddings to work but will miss the key.
Since the PR explicitly introduces xAI support, both example files should document the full set of required env vars (or intentionally omit both).
Summary of changes
What changed
Environment / docs
- Added xAI-related variables to example env files:
.env.example:XAI_API_KEY,XAI_EMBEDDING_MODEL(defaults togrok-v1-embedding), andOPENAI_API_KEY..env.local.example: addedXAI_EMBEDDING_MODELunder an xAI Embeddings section.
Embedding generation
lib/actions/chat.ts- Added
extractTextForEmbedding()to extract onlytextparts from multimodal messagecontent. - Updated embedding generation to call
embedMany()with extracted text values. - Refactored message persistence mapping to store JSON-stringified content (
jsonContent) while embedding is computed from extracted text. - Added
id: crypto.randomUUID()for the drawing-context “data” message.
- Added
Provider selection
lib/utils/index.ts- Enhanced
getEmbeddingModel()to optionally use xAI embeddings whenXAI_API_KEYandXAI_EMBEDDING_MODELare provided. - Falls back to OpenAI
text-embedding-3-smallwhen xAI is not configured. - Updated warning message when embeddings are unavailable.
- Enhanced
| function extractTextForEmbedding(content: any): string { | ||
| if (typeof content === 'string') { | ||
| return content; | ||
| } | ||
| if (Array.isArray(content)) { | ||
| return content | ||
| .filter((part: any) => part.type === 'text') | ||
| .map((part: any) => part.text) | ||
| .join('\n'); | ||
| } | ||
| return ''; | ||
| } | ||
|
|
There was a problem hiding this comment.
extractTextForEmbedding(content: any) introduces any and silently drops non-string / non-array content. This is both an unsafe-but-type-valid pattern and can mask bugs where content shape changes (e.g., object-based content that isn’t an array of parts). It also assumes parts have { type, text } without guarding text being a string.
Given this is security/quality sensitive (it affects what gets embedded and searched), tighten the input type and add a conservative fallback for known shapes rather than returning an empty string for everything else.
Suggestion
Prefer typing to the shapes you actually support (string or Vercel AI-style parts) and make the fallback explicit. For example:
type TextPart = { type: 'text'; text: string };
type ContentPart = TextPart | { type: string; [key: string]: unknown };
function extractTextForEmbedding(content: unknown): string {
if (typeof content === 'string') return content;
if (Array.isArray(content)) {
return content
.filter((p): p is TextPart =>
!!p && typeof p === 'object' && (p as any).type === 'text' && typeof (p as any).text === 'string'
)
.map(p => p.text)
.join('\n');
}
// If content is an object (unexpected), consider JSON-stringifying *only if* you really want it embedded.
return '';
}This removes any, hardens extraction, and makes it harder to accidentally embed garbage or nothing. Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this change.
| const messagesNeedingEmbeddings = messagesWithExtractedText.filter(m => | ||
| (m.role === 'user' || m.role === 'assistant') && m.extractedText && m.extractedText.length > 0 | ||
| ); | ||
|
|
||
| let embeddings: number[][] = []; | ||
| if (messagesNeedingEmbeddings.length > 0) { | ||
| const { embeddings } = await embedMany({ | ||
| const result = await embedMany({ | ||
| model: embeddingModel, | ||
| values: messagesNeedingEmbeddings.map(m => m.content as string), | ||
| values: messagesNeedingEmbeddings.map(m => m.extractedText), | ||
| }); | ||
| embeddings = result.embeddings; | ||
| } | ||
|
|
||
| // Map embeddings back to messages | ||
| messagesToProcess.forEach(m => { | ||
| const index = messagesNeedingEmbeddings.findIndex(nm => nm.id === m.id); | ||
| processedMessages.push({ | ||
| ...m, | ||
| chatId: chat.id, | ||
| embedding: index !== -1 ? embeddings[index] : null | ||
| }); | ||
| }); | ||
| } else { | ||
| messagesToProcess.forEach(m => { | ||
| processedMessages.push({ ...m, chatId: chat.id, embedding: null }); | ||
| // Map embeddings back to messages | ||
| messagesWithExtractedText.forEach(m => { | ||
| const index = messagesNeedingEmbeddings.findIndex(nm => nm.id === m.id); | ||
| processedMessages.push({ | ||
| id: m.id, | ||
| chatId: chat.id, | ||
| userId: effectiveUserId, | ||
| role: m.role, | ||
| content: m.jsonContent, | ||
| createdAt: m.createdAt ? new Date(m.createdAt) : new Date(), | ||
| embedding: index !== -1 ? embeddings[index] : null | ||
| }); | ||
| } | ||
| }); |
There was a problem hiding this comment.
The embedding mapping uses findIndex inside a forEach, making the mapping O(n²) for the number of messages. For long chats, this can become a noticeable server-side slowdown.
This is easy to make O(n) by building a lookup map from message id → embedding index (or directly id → embedding) once.
Suggestion
Precompute a map for constant-time lookups:
const embeddingById = new Map<string, number[]>();
messagesNeedingEmbeddings.forEach((m, i) => {
const e = embeddings[i];
if (e) embeddingById.set(m.id, e);
});
messagesWithExtractedText.forEach(m => {
processedMessages.push({
id: m.id,
chatId: chat.id,
userId: effectiveUserId,
role: m.role,
content: m.jsonContent,
createdAt: m.createdAt ? new Date(m.createdAt) : new Date(),
embedding: embeddingById.get(m.id) ?? null,
});
});Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this optimization.
| export function getEmbeddingModel() { | ||
| const openaiApiKey = process.env.OPENAI_API_KEY; | ||
| const xaiApiKey = process.env.XAI_API_KEY; | ||
| const xaiEmbeddingModel = process.env.XAI_EMBEDDING_MODEL; | ||
|
|
||
| // Use xAI if both key and model are explicitly provided | ||
| if (xaiApiKey && xaiEmbeddingModel) { | ||
| const xai = createOpenAI({ | ||
| apiKey: xaiApiKey, | ||
| baseURL: 'https://api.x.ai/v1', | ||
| }); | ||
| return xai.embedding(xaiEmbeddingModel); | ||
| } | ||
|
|
||
| // Fallback to OpenAI text-embedding-3-small | ||
| if (!openaiApiKey) { | ||
| console.warn('OPENAI_API_KEY is not set. Embedding functionality will be unavailable.'); | ||
| console.warn('Neither XAI_EMBEDDING_MODEL nor OPENAI_API_KEY is set. Embedding functionality will be unavailable.'); | ||
| return null; | ||
| } |
There was a problem hiding this comment.
getEmbeddingModel() now warns: Neither XAI_EMBEDDING_MODEL nor OPENAI_API_KEY is set..., but the actual requirements for xAI are XAI_API_KEY and XAI_EMBEDDING_MODEL. As written, a user could set XAI_EMBEDDING_MODEL without XAI_API_KEY and get a misleading warning implying OpenAI key is the only missing piece.
This makes operational troubleshooting harder.
Suggestion
Adjust the warning to reflect the real configuration matrix:
if (!openaiApiKey && !(xaiApiKey && xaiEmbeddingModel)) {
console.warn(
'Embeddings unavailable. Configure either (XAI_API_KEY + XAI_EMBEDDING_MODEL) or OPENAI_API_KEY.'
);
return null;
}Reply with "@CharlieHelps yes please" if you'd like me to add a commit with this improved warning logic.
2 participants
Improved the embedding generation process by extracting only the text content from multimodal messages, ensuring better search and RAG quality. Also updated the embedding model configuration to support xAI as an optional provider while maintaining OpenAI's text-embedding-3-small as the functional default. Documentation for new environment variables was added to example files.
PR created automatically by Jules for task 18237254138839941279 started by @ngoiyaeric