Improve RAG context injection #223
Conversation
WalkthroughThis update introduces chunk context awareness and expansion in document processing and retrieval. It adds a Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Retriever
participant VectorStore
User->>Retriever: Ask question
Retriever->>VectorStore: Vector search with prompt
VectorStore-->>Retriever: Return top N chunks
Retriever->>Retriever: Group chunks by path
loop For each path group
Retriever->>VectorStore: Fetch neighboring chunks (expand context)
VectorStore-->>Retriever: Return expanded chunks
Retriever->>Retriever: Concatenate and sort chunks, preserve relevance
end
Retriever-->>User: Return expanded relevant documents
Suggested reviewers
Poem
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
npm error Exit handler never called! 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
Deploying with
|
| Status | Name | Latest Commit | Preview URL | Updated (UTC) |
|---|---|---|---|---|
| ✅ Deployment successful! View logs |
docs | 5f8c376 | Commit Preview URL | Jul 10 2025, 12:10 AM |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (1)
agent-docs/src/agents/doc-processing/docs-orchestrator.ts (1)
92-94: Use the correct logger specifier for the upsert countWe’ve confirmed via the JavaScript SDK docs that
ctx.vector.upsertreturns aPromise<string[]>—i.e., an array of string IDs—so accessingresult.lengthis safe. To improve clarity and formatting, switch the logger’s placeholder from%o(object) to%d(decimal number).• File: agent-docs/src/agents/doc-processing/docs-orchestrator.ts
Line 93Suggested change:
- ctx.logger.info('Upserted chunk: %o', result.length); + ctx.logger.info('Upserted chunk count: %d', result.length);
🧹 Nitpick comments (3)
agent-docs/src/agents/doc-processing/docs-orchestrator.ts (1)
96-96: Use correct format specifier for numeric value.- ctx.logger.info('Upserted total %o chunks for file: %s', chunks.length, logicalPath); + ctx.logger.info('Upserted total %d chunks for file: %s', chunks.length, logicalPath);agent-docs/src/agents/doc-qa/retriever.ts (2)
22-22: Consider making contextWindow configurable.The context window is hardcoded to 1. This should be configurable to allow flexibility in how much surrounding context is included.
Consider moving this to a configuration constant or accepting it as a parameter:
async function expandPathGroup( ctx: AgentContext, path: string, pathChunks: Array<{ path: string; content: string; relevanceScore?: number; chunkIndex?: number; - }> + }>, + contextWindow: number = 1 ): Promise<RelevantDoc | null> { - const contextWindow = 1; // Get 1 chunk before and after each chunk
92-92: Use appropriate format specifier for relevance score.- ctx.logger.debug('Expanded path %s with %d chunks (range: %s) score %d', path, foundChunks.length, chunkRange, bestScore); + ctx.logger.debug('Expanded path %s with %d chunks (range: %s) score %f', path, foundChunks.length, chunkRange, bestScore);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
agent-docs/config.ts(1 hunks)agent-docs/src/agents/doc-processing/chunk-mdx.ts(2 hunks)agent-docs/src/agents/doc-processing/docs-orchestrator.ts(2 hunks)agent-docs/src/agents/doc-processing/docs-processor.ts(1 hunks)agent-docs/src/agents/doc-processing/types.ts(1 hunks)agent-docs/src/agents/doc-qa/rag.ts(3 hunks)agent-docs/src/agents/doc-qa/retriever.ts(1 hunks)agent-docs/src/agents/doc-qa/types.ts(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
`agent-docs/src/agents/**/*.ts`: Use TypeScript for better type safety and IDE s...
agent-docs/src/agents/**/*.ts: Use TypeScript for better type safety and IDE support
Import types from@agentuity/sdk
Use structured error handling with try/catch blocks
Leverage the provided logger for consistent logging
Use the storage APIs for persisting data
📄 Source: CodeRabbit Inference Engine (agent-docs/.cursor/rules/sdk.mdc)
List of files the instruction was applied to:
agent-docs/src/agents/doc-processing/docs-processor.tsagent-docs/src/agents/doc-processing/types.tsagent-docs/src/agents/doc-processing/chunk-mdx.tsagent-docs/src/agents/doc-qa/types.tsagent-docs/src/agents/doc-processing/docs-orchestrator.tsagent-docs/src/agents/doc-qa/rag.tsagent-docs/src/agents/doc-qa/retriever.ts
🧠 Learnings (8)
📓 Common learnings
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:7-15
Timestamp: 2025-06-18T04:07:57.583Z
Learning: In the agent-docs project, the ChunkMetadata type requires title and description fields to be mandatory strings rather than optional, indicating that the system ensures these fields are always populated during document processing.
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:27-33
Timestamp: 2025-06-18T04:09:27.663Z
Learning: The `embedChunks` function in the agent-docs project maintains strict order and count consistency with input chunks, so positional correspondence between chunks.map() and embeddings[index] is reliable without additional safeguarding.
agent-docs/src/agents/doc-processing/docs-processor.ts (2)
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:27-33
Timestamp: 2025-06-18T04:09:27.663Z
Learning: The `embedChunks` function in the agent-docs project maintains strict order and count consistency with input chunks, so positional correspondence between chunks.map() and embeddings[index] is reliable without additional safeguarding.
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:7-15
Timestamp: 2025-06-18T04:07:57.583Z
Learning: In the agent-docs project, the ChunkMetadata type requires title and description fields to be mandatory strings rather than optional, indicating that the system ensures these fields are always populated during document processing.
agent-docs/src/agents/doc-processing/types.ts (5)
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:7-15
Timestamp: 2025-06-18T04:07:57.583Z
Learning: In the agent-docs project, the ChunkMetadata type requires title and description fields to be mandatory strings rather than optional, indicating that the system ensures these fields are always populated during document processing.
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:27-33
Timestamp: 2025-06-18T04:09:27.663Z
Learning: The `embedChunks` function in the agent-docs project maintains strict order and count consistency with input chunks, so positional correspondence between chunks.map() and embeddings[index] is reliable without additional safeguarding.
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : All code should be in Typescript format
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Import types from `@agentuity/sdk`
agent-docs/config.ts (3)
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : Prefer naming the default function Agent or the name of the Agent based on the context of the Agent description
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : All code should be in Typescript format
agent-docs/src/agents/doc-processing/chunk-mdx.ts (4)
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:7-15
Timestamp: 2025-06-18T04:07:57.583Z
Learning: In the agent-docs project, the ChunkMetadata type requires title and description fields to be mandatory strings rather than optional, indicating that the system ensures these fields are always populated during document processing.
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:27-33
Timestamp: 2025-06-18T04:09:27.663Z
Learning: The `embedChunks` function in the agent-docs project maintains strict order and count consistency with input chunks, so positional correspondence between chunks.map() and embeddings[index] is reliable without additional safeguarding.
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : All code should be in Typescript format
agent-docs/src/agents/doc-qa/types.ts (5)
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:7-15
Timestamp: 2025-06-18T04:07:57.583Z
Learning: In the agent-docs project, the ChunkMetadata type requires title and description fields to be mandatory strings rather than optional, indicating that the system ensures these fields are always populated during document processing.
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Import types from `@agentuity/sdk`
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : All code should be in Typescript format
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : Prefer loading types from the node modules package `@agentuity/sdk` in the node_modules folder
agent-docs/src/agents/doc-qa/rag.ts (5)
Learnt from: afterrburn
PR: agentuity/docs#208
File: agent-docs/src/agents/doc-qa/index.ts:27-28
Timestamp: 2025-06-28T18:06:02.340Z
Learning: The answerQuestion function in agent-docs/src/agents/doc-qa/rag.ts handles errors internally by logging them and returning a fallback JSON answer, so it doesn't require external try-catch blocks at the call site.
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : Prefer naming the default function Agent or the name of the Agent based on the context of the Agent description
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Prefer using the `agentuity agent create` command to create a new Agent
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agentuity.mdc:0-0
Timestamp: 2025-07-01T12:36:39.469Z
Learning: Applies to agent-docs/**/agentuity.yaml : Do not suggest edits to the Agentuity AI Configuration file (agentuity.yaml)
agent-docs/src/agents/doc-qa/retriever.ts (8)
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : Prefer naming the default function Agent or the name of the Agent based on the context of the Agent description
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use TypeScript for better type safety and IDE support
Learnt from: afterrburn
PR: agentuity/docs#195
File: agent-docs/src/agents/doc-processing/docs-processor.ts:27-33
Timestamp: 2025-06-18T04:09:27.663Z
Learning: The `embedChunks` function in the agent-docs project maintains strict order and count consistency with input chunks, so positional correspondence between chunks.map() and embeddings[index] is reliable without additional safeguarding.
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : All code should be in Typescript format
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : Prefer loading types from the node modules package `@agentuity/sdk` in the node_modules folder
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Import types from `@agentuity/sdk`
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/agent.mdc:0-0
Timestamp: 2025-07-01T12:36:36.319Z
Learning: Applies to agent-docs/**/src/agents/**/index.ts : The file should export a default function
Learnt from: CR
PR: agentuity/docs#0
File: agent-docs/.cursor/rules/sdk.mdc:0-0
Timestamp: 2025-07-01T12:36:46.237Z
Learning: Applies to agent-docs/src/agents/**/*.ts : Use structured error handling with try/catch blocks
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: Workers Builds: docs
🔇 Additional comments (10)
agent-docs/config.ts (1)
2-2: LGTM: Vector search reduction aligns with improved RAG approach.Reducing the vector search number from 20 to 10 makes sense given the new chunk grouping and expansion strategy. This should improve performance while maintaining relevance quality.
agent-docs/src/agents/doc-processing/chunk-mdx.ts (2)
14-14: LGTM: Added totalChunks property for chunk context awareness.The addition of
totalChunksto the Chunk type provides valuable context about the document's chunk distribution, supporting the improved RAG approach.
148-148: LGTM: Correct implementation of totalChunks assignment.The assignment of
totalChunks: chunks.lengthis accurate and provides each chunk with awareness of the total number of chunks in the document.agent-docs/src/agents/doc-processing/docs-processor.ts (1)
27-27: LGTM: Proper propagation of totalChunks to vector metadata.The addition of
totalChunks: chunk.totalChunksto the metadata object ensures consistent chunk context information throughout the embedding pipeline.agent-docs/src/agents/doc-processing/types.ts (1)
22-22: LGTM: Type consistency maintained with totalChunks addition.The addition of
totalChunks: numberto the ChunkMetadata type maintains proper type safety and consistency with the extended Chunk type.agent-docs/src/agents/doc-qa/types.ts (1)
6-8: LGTM: Schema extensions support enhanced RAG context injection.The addition of
relevanceScore,chunkRange, andchunkIndexas optional fields properly extends the schema to support the new chunk grouping and expansion approach. The optional nature and correct typing make these additions flexible and type-safe.agent-docs/src/agents/doc-processing/docs-orchestrator.ts (1)
15-19: Clarify the use of whitespace query for vector search.Using a single space
' 'as the query appears to be a workaround to retrieve all vectors. Is this the intended API pattern? Consider documenting this behavior or checking if there's a more explicit way to query all vectors with the given metadata filter.agent-docs/src/agents/doc-qa/rag.ts (1)
17-101: System prompt refactoring looks good!The simplified prompt structure with a single
<USER_QUESTION>tag and the addition of user context (software engineer) should help generate more targeted responses. The separation of system prompt and user-specific instructions is clean.agent-docs/src/agents/doc-qa/retriever.ts (2)
173-184: Well-structured chunk expansion implementation!The approach of grouping chunks by path before expansion is efficient and avoids redundant queries. The error handling and logging throughout the function provide good observability.
142-145: Verify and strengthen the vector search return type forsimilarityWe’re currently doing:
const relevanceScore = (vector as any).similarity;which sidesteps TypeScript’s checks. Since we couldn’t locate a local
Vector…Result/Responsetype in the repo, please:
- Inspect the
ctx.vector.searchreturn type in@agentuity/sdk(via its.d.ts).- Import or extend the correct interface (e.g.
VectorSearchResult) so it includessimilarity: number.- Replace the
as anycast with that concrete type.Example diff in agent-docs/src/agents/doc-qa/retriever.ts:
-import type { AgentContext } from '@agentuity/sdk'; +import type { AgentContext, VectorSearchResult } from '@agentuity/sdk'; … -for (const vector of vectors) { +for (const vector of vectors as VectorSearchResult[]) { // … - const relevanceScore = (vector as any).similarity; + const relevanceScore = vector.similarity; // … }If the SDK doesn’t yet export such a type, declare a local interface that mirrors its shape and use that instead.
| // Delete vectors one by one to avoid issues with large batches | ||
| for (const vector of vectors) { | ||
| await ctx.vector.delete(vectorStoreName, vector.key); | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Consider batch deletion for better performance.
Deleting vectors one by one could be inefficient, especially with 100 vectors per batch. This results in 100 individual API calls per iteration.
Check if the vector API supports batch deletion:
If batch deletion is available, consider refactoring:
- // Delete vectors one by one to avoid issues with large batches
- for (const vector of vectors) {
- await ctx.vector.delete(vectorStoreName, vector.key);
- }
+ // Delete vectors in batch
+ const keys = vectors.map(v => v.key);
+ await ctx.vector.deleteMany(vectorStoreName, keys);🏁 Script executed:
#!/bin/bash
# Search for batch delete methods in the SDK or vector-related files
rg -A 5 "batch.*delete|delete.*batch|deleteMany|bulkDelete" --type tsLength of output: 335
Unify and parallelize vector deletions for improved performance and reliability
No native batch‐delete method was found in the SDK, so we should at least run deletions in parallel and add error handling. Replace the sequential for…await loop with something like:
- // Delete vectors one by one to avoid issues with large batches
- for (const vector of vectors) {
- await ctx.vector.delete(vectorStoreName, vector.key);
- }
+ // Delete vectors in parallel with individual error handling
+ const results = await Promise.allSettled(
+ vectors.map(v => ctx.vector.delete(vectorStoreName, v.key))
+ );
+ results.forEach((res, i) => {
+ if (res.status === 'rejected') {
+ logger.error(
+ `Failed to delete vector ${vectors[i].key}:`,
+ res.reason
+ );
+ }
+ });• This runs all deletions concurrently to reduce overall latency.
• Using Promise.allSettled lets us log failures without aborting the entire batch.
• For very large batches, consider chunking the vectors array to limit in‐flight requests.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Delete vectors one by one to avoid issues with large batches | |
| for (const vector of vectors) { | |
| await ctx.vector.delete(vectorStoreName, vector.key); | |
| } | |
| // Delete vectors in parallel with individual error handling | |
| const results = await Promise.allSettled( | |
| vectors.map(v => ctx.vector.delete(vectorStoreName, v.key)) | |
| ); | |
| results.forEach((res, i) => { | |
| if (res.status === 'rejected') { | |
| logger.error( | |
| `Failed to delete vector ${vectors[i].key}:`, | |
| res.reason | |
| ); | |
| } | |
| }); |
🤖 Prompt for AI Agents
In agent-docs/src/agents/doc-processing/docs-orchestrator.ts around lines 25 to
28, the current code deletes vectors sequentially which is inefficient. Replace
the for-await loop with a parallel deletion approach using Promise.allSettled to
run all deletions concurrently and handle errors individually by logging
failures without stopping the entire batch. Optionally, implement chunking of
the vectors array if the batch size is very large to limit the number of
simultaneous requests.
| const expandedContent = sortedChunks.map(chunk => chunk.text).join('\n\n'); | ||
|
|
||
| // Find the best relevance score from the original chunks | ||
| const bestScore = Math.max(...pathChunks.map(chunk => chunk.relevanceScore || 0)); |
There was a problem hiding this comment.
Handle empty array case for Math.max.
Math.max() returns -Infinity when called with an empty array or when all relevanceScores are undefined.
- const bestScore = Math.max(...pathChunks.map(chunk => chunk.relevanceScore || 0));
+ const scores = pathChunks.map(chunk => chunk.relevanceScore || 0);
+ const bestScore = scores.length > 0 ? Math.max(...scores) : 0;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| const bestScore = Math.max(...pathChunks.map(chunk => chunk.relevanceScore || 0)); | |
| const scores = pathChunks.map(chunk => chunk.relevanceScore || 0); | |
| const bestScore = scores.length > 0 ? Math.max(...scores) : 0; |
🤖 Prompt for AI Agents
In agent-docs/src/agents/doc-qa/retriever.ts at line 85, the use of Math.max on
an array derived from pathChunks can result in -Infinity if the array is empty
or all relevanceScores are undefined. To fix this, check if pathChunks is empty
or if all relevanceScores are undefined before calling Math.max, and provide a
default value (e.g., 0) to use in that case to avoid -Infinity.
| for (const chunk of pathChunks) { | ||
| if (chunk.chunkIndex !== undefined) { | ||
| const targetIndex = chunk.chunkIndex; | ||
|
|
||
| // Add the chunk itself and its neighbors | ||
| expandedChunkIndices.add(targetIndex - contextWindow); | ||
| expandedChunkIndices.add(targetIndex); | ||
| expandedChunkIndices.add(targetIndex + contextWindow); | ||
| } | ||
| try { | ||
| const vectors = await ctx.vector.search(VECTOR_STORE_NAME, dbQuery); | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Add upper bound validation for chunk indices.
The code adds neighboring chunk indices without checking against totalChunks. While negative indices are filtered out, indices beyond the document's total chunks are not validated.
First, let's verify if totalChunks metadata is available:
If available, add upper bound checking:
// Add neighbors for each chunk to the set
for (const chunk of pathChunks) {
if (chunk.chunkIndex !== undefined) {
const targetIndex = chunk.chunkIndex;
+ // Assuming totalChunks is available in metadata
+ const totalChunks = chunk.totalChunks || Number.MAX_SAFE_INTEGER;
// Add the chunk itself and its neighbors
- expandedChunkIndices.add(targetIndex - contextWindow);
+ if (targetIndex - contextWindow >= 0) {
+ expandedChunkIndices.add(targetIndex - contextWindow);
+ }
expandedChunkIndices.add(targetIndex);
- expandedChunkIndices.add(targetIndex + contextWindow);
+ if (targetIndex + contextWindow < totalChunks) {
+ expandedChunkIndices.add(targetIndex + contextWindow);
+ }
}
}🏁 Script executed:
#!/bin/bash
# Check if totalChunks is part of the chunk metadata
rg -A 5 "totalChunks" --type tsLength of output: 1901
Ensure upper‐bound validation when expanding chunk indices
Add checks against chunk.totalChunks so neighbor indices never exceed the last chunk.
- File:
agent-docs/src/agents/doc-qa/retriever.ts - Loop over
pathChunksat lines ~26–35
Suggested diff:
for (const chunk of pathChunks) {
if (chunk.chunkIndex !== undefined) {
const targetIndex = chunk.chunkIndex;
+ const totalChunks = chunk.totalChunks;
- expandedChunkIndices.add(targetIndex - contextWindow);
+ if (targetIndex - contextWindow >= 0) {
+ expandedChunkIndices.add(targetIndex - contextWindow);
+ }
expandedChunkIndices.add(targetIndex);
- expandedChunkIndices.add(targetIndex + contextWindow);
+ if (targetIndex + contextWindow < totalChunks) {
+ expandedChunkIndices.add(targetIndex + contextWindow);
+ }
}
}This guarantees all added indices stay within [0, totalChunks – 1].
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for (const chunk of pathChunks) { | |
| if (chunk.chunkIndex !== undefined) { | |
| const targetIndex = chunk.chunkIndex; | |
| // Add the chunk itself and its neighbors | |
| expandedChunkIndices.add(targetIndex - contextWindow); | |
| expandedChunkIndices.add(targetIndex); | |
| expandedChunkIndices.add(targetIndex + contextWindow); | |
| } | |
| try { | |
| const vectors = await ctx.vector.search(VECTOR_STORE_NAME, dbQuery); | |
| } | |
| for (const chunk of pathChunks) { | |
| if (chunk.chunkIndex !== undefined) { | |
| const targetIndex = chunk.chunkIndex; | |
| const totalChunks = chunk.totalChunks; | |
| // Add the chunk itself and its neighbors | |
| if (targetIndex - contextWindow >= 0) { | |
| expandedChunkIndices.add(targetIndex - contextWindow); | |
| } | |
| expandedChunkIndices.add(targetIndex); | |
| if (targetIndex + contextWindow < totalChunks) { | |
| expandedChunkIndices.add(targetIndex + contextWindow); | |
| } | |
| } | |
| } |
🤖 Prompt for AI Agents
In agent-docs/src/agents/doc-qa/retriever.ts around lines 26 to 35, the code
expands chunk indices by adding neighbors without checking if these indices
exceed the upper bound defined by chunk.totalChunks. To fix this, add validation
to ensure that when adding targetIndex - contextWindow and targetIndex +
contextWindow, the indices stay within the range 0 to totalChunks - 1,
preventing out-of-bounds errors.
agent-docs/src/agents/doc-qa/rag.ts
Outdated
| === CONTEXT === | ||
| You will receive both the user's ORIGINAL question and a REPHRASED version that was optimized for document search. The rephrased version helped find the relevant documents, but you should answer the user's original intent. | ||
| Your role is to be as helpful as possible. You must first |
|
@afterrburn the SDK default if you don't pass a limit in your query is 10. you can pass a |
Signed-off-by: Seng Rith <50646727+afterrburn@users.noreply.github.com>
Since the SDK for vector search is currently broken (only returning a maximum of 10 results), this is a good opportunity to revise our approach. Previously, we retrieved the full document associated with a chunk, which provided rich context but often led to saturated and slower responses.
The new strategy retrieves only the content surrounding relevant chunks, providing more targeted context to the LLM. This avoids unnecessary overhead—some markdowns have up to 75 chunks, but only a few are needed to answer most questions.
I also reduced the number of relevant vector results from 20 to 10 to match the current SDK behavior. We can increase this when it makes sense.
Summary by CodeRabbit
New Features
Improvements