Improve RAG context injection

agent-docs/config.ts

@@ -1,2 +1,2 @@
 export const VECTOR_STORE_NAME = process.env.VECTOR_STORE_NAME || 'docs';
-export const vectorSearchNumber = 20;
+export const vectorSearchNumber = 10;

agent-docs/src/agents/doc-processing/chunk-mdx.ts

@@ -11,6 +11,7 @@ import matter from 'gray-matter';
 export type Chunk = {
   id: string;
   chunkIndex: number;
+  totalChunks: number;
   contentType: string;
   heading: string;
   text: string;
@@ -144,6 +145,7 @@ export async function chunkAndEnrichDoc(fileContent: string): Promise<Chunk[]> {
     return {
       id: crypto.randomUUID(),
       chunkIndex: idx,
+      totalChunks: chunks.length,
       contentType: chunk.metadata.contentType,
       heading: currentHeading,
       text: chunk.pageContent,

agent-docs/src/agents/doc-processing/docs-orchestrator.ts

@@ -8,18 +8,31 @@ import type { SyncPayload, SyncStats } from './types';
  */
 async function removeVectorsByPath(ctx: AgentContext, logicalPath: string, vectorStoreName: string) {
   ctx.logger.info('Removing vectors for path: %s', logicalPath);
-  const vectors = await ctx.vector.search(vectorStoreName, {
-    query: ' ',
-    limit: 10000,
-    metadata: { path: logicalPath },
-  });
+
+  let totalDeleted = 0;
+
+  while (true) {
+    const vectors = await ctx.vector.search(vectorStoreName, {
+      query: ' ',
+      limit: 100,
+      metadata: { path: logicalPath },
+    });
+
+    if (!Array.isArray(vectors) || vectors.length === 0) {
+      break;
+    }
 
-  if (vectors.length > 0) {
     // Delete vectors one by one to avoid issues with large batches
     for (const vector of vectors) {
       await ctx.vector.delete(vectorStoreName, vector.key);
     }
-    ctx.logger.info('Removed %d vectors for path: %s', vectors.length, logicalPath);
+
+    totalDeleted += vectors.length;
+    ctx.logger.info('Deleted %d vectors (total: %d) for path: %s', vectors.length, totalDeleted, logicalPath);
+  }
+
+  if (totalDeleted > 0) {
+    ctx.logger.info('Completed removal of %d vectors for path: %s', totalDeleted, logicalPath);
   } else {
     ctx.logger.info('No vectors found for path: %s', logicalPath);
   }
@@ -80,6 +93,8 @@ export async function syncDocsFromPayload(ctx: AgentContext, payload: SyncPayloa
         ctx.logger.info('Upserted chunk: %o', result.length);
       }
 
+      ctx.logger.info('Upserted total %o chunks for file: %s', chunks.length, logicalPath);
+
       processed++;
       ctx.logger.info('Successfully processed file: %s (%d chunks)', logicalPath, chunks.length);
     } catch (err) {

agent-docs/src/agents/doc-processing/docs-processor.ts

@@ -24,6 +24,7 @@ async function createVectorEmbedding(chunks: Chunk[]): Promise<VectorUpsertParam
     }
     const metadata: ChunkMetadata = {
       chunkIndex: chunk.chunkIndex,
+      totalChunks: chunk.totalChunks,
       contentType: chunk.contentType,
       heading: chunk.heading,
       title: chunk.title,

agent-docs/src/agents/doc-processing/types.ts

@@ -19,6 +19,7 @@ export interface SyncStats {
 
 export type ChunkMetadata = {
   chunkIndex: number;
+  totalChunks: number;
   contentType: string;
   heading: string;
   title: string;

agent-docs/src/agents/doc-qa/rag.ts

@@ -18,12 +18,13 @@ export default async function answerQuestion(ctx: AgentContext, prompt: string)
 You are Agentuity's developer-documentation assistant.
 
 === 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 and try to assist user by answering their questions.
 
 === RULES ===
 1. Use ONLY the content inside <DOCS> tags to craft your reply. If the required information is missing, state that the docs do not cover it.
 2. Never fabricate or guess undocumented details.
-3. Focus on answering the ORIGINAL QUESTION, using the documents found via the rephrased search.
+3. Focus on answering the QUESTION with the available <DOCS> provided to you. Keep in mind some <DOCS> might not be relevant,
+   so pick the ones that is relevant to the user's question.
 4. Ambiguity handling:
    • When <DOCS> contains more than one distinct workflow or context that could satisfy the question, do **not** choose for the user.
    • Briefly (≤ 2 sentences each) summarise each plausible interpretation and ask **one** clarifying question so the user can pick a path.
@@ -83,13 +84,9 @@ agentuity agent create [name] [description] [auth_type]
 
 > **Note**: This command will create the agent in the Agentuity Cloud and set up local files.
 
-<ORIGINAL_QUESTION>
-${prompt}
-</ORIGINAL_QUESTION>
-
-<REPHRASED_SEARCH_QUERY>
+<USER_QUESTION>
 ${rephrasedPrompt}
-</REPHRASED_SEARCH_QUERY>
+</USER_QUESTION>
 
 <DOCS>
 ${JSON.stringify(relevantDocs, null, 2)}
@@ -100,7 +97,7 @@ ${JSON.stringify(relevantDocs, null, 2)}
         const result = await generateObject({
             model: openai('gpt-4o'),
             system: systemPrompt,
-            prompt: `Please answer the original question using the documentation found via the rephrased search query. Your answer should cater toward the original user prompt rather than the rephrased version of the query.`,
+            prompt: `The user is mostly a software engineer. Your answer should be concise, straightforward and in most cases, supplying the answer with examples code snipped is ideal.`,
             schema: AnswerSchema,
             maxTokens: 2048,
         });

agent-docs/src/agents/doc-qa/retriever.ts

@@ -1,80 +1,190 @@
 import type { AgentContext } from '@agentuity/sdk';
 
-import type { ChunkMetadata } from '../doc-processing/types';
 import { VECTOR_STORE_NAME, vectorSearchNumber } from '../../../config';
 import type { RelevantDoc } from './types';
 
-export async function retrieveRelevantDocs(ctx: AgentContext, prompt: string): Promise<RelevantDoc[]> {
-    const dbQuery = {
-      query: prompt,
-      limit: vectorSearchNumber
+
+
+/**
+ * Expands a group of chunks from the same path by creating a set of all needed chunk indices
+ * and querying for them once
+ */
+async function expandPathGroup(
+  ctx: AgentContext,
+  path: string,
+  pathChunks: Array<{
+    path: string;
+    content: string;
+    relevanceScore?: number;
+    chunkIndex?: number;
+  }>
+): Promise<RelevantDoc | null> {
+  const contextWindow = 1; // Get 1 chunk before and after each chunk
+  const expandedChunkIndices = new Set<number>();
+
+  // Add neighbors for each chunk to the set
+  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);
+  }
 
-      const uniquePaths = new Set<string>();
+  // Remove negative indices
+  const validIndices = Array.from(expandedChunkIndices).filter(index => index >= 0);
 
-      vectors.forEach(vec => {
-        if (!vec.metadata) {
-          ctx.logger.warn('Vector missing metadata');
-          return;
-        }
-        const path = typeof vec.metadata.path === 'string' ? vec.metadata.path : undefined;
-        if (!path) {
-          ctx.logger.warn('Vector metadata path is not a string');
-          return;
-        }
-        uniquePaths.add(path);
-      });
+  if (validIndices.length === 0) {
+    ctx.logger.warn('No valid chunk indices found for path: %s', path);
+    return null;
+  }
 
-      const docs = await Promise.all(
-        Array.from(uniquePaths).map(async path => ({
-          path,
-          content: await retrieveDocumentBasedOnPath(ctx, path)
-        }))
-      );
+  // Sort indices
+  validIndices.sort((a, b) => a - b);
 
-      return docs;
-    } catch (err) {
-      ctx.logger.error('Error retrieving relevant docs: %o', err);
-      return [];
+  try {
+    // Query for all chunks at once
+    const chunkQueries = validIndices.map(index =>
+      ctx.vector.search(VECTOR_STORE_NAME, {
+        query: ' ',
+        limit: 1,
+        metadata: { path: path, chunkIndex: index }
+      })
+    );
+
+    const results = await Promise.all(chunkQueries);
+
+    // Collect found chunks
+    const foundChunks: Array<{ index: number; text: string }> = [];
+
+    for (const result of results) {
+      if (result.length > 0 && result[0] && result[0].metadata) {
+        const metadata = result[0].metadata;
+        if (typeof metadata.chunkIndex === 'number' && typeof metadata.text === 'string') {
+          foundChunks.push({
+            index: metadata.chunkIndex,
+            text: metadata.text
+          });
+        }
+      }
+    }
+
+    if (foundChunks.length === 0) {
+      ctx.logger.warn('No chunks found for path: %s', path);
+      return null;
     }
+
+    // Sort by index and combine content
+    const sortedChunks = foundChunks.sort((a, b) => a.index - b.index);
+    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));
+
+    // Create chunk range
+    const minIndex = Math.min(...sortedChunks.map(c => c.index));
+    const maxIndex = Math.max(...sortedChunks.map(c => c.index));
+    const chunkRange = minIndex === maxIndex ? `${minIndex}` : `${minIndex}-${maxIndex}`;
+
+    ctx.logger.debug('Expanded path %s with %d chunks (range: %s) score %d', path, foundChunks.length, chunkRange, bestScore);
+
+    return {
+      path,
+      content: expandedContent,
+      relevanceScore: bestScore,
+      chunkRange,
+      chunkIndex: undefined // Not applicable for grouped chunks
+    };
+
+  } catch (err) {
+    ctx.logger.error('Error expanding path group %s: %o', path, err);
+    return null;
   }
-
-  async function retrieveDocumentBasedOnPath(ctx: AgentContext, path: string): Promise<string> {
-    const dbQuery = {
-      query: ' ',
-      limit: 1000,
-      metadata: {
-        path: path
+}
+
+export async function retrieveRelevantDocs(ctx: AgentContext, prompt: string): Promise<RelevantDoc[]> {
+  const dbQuery = {
+    query: prompt,
+    limit: vectorSearchNumber
+  }
+  try {
+    const vectors = await ctx.vector.search(VECTOR_STORE_NAME, dbQuery);
+
+    ctx.logger.debug('Vector search returned %d results. First vector structure: %o', 
+                     vectors.length, vectors[0]);
+
+    // Process each relevant chunk and expand with context
+    const relevantChunks: Array<{
+      path: string;
+      content: string;
+      relevanceScore?: number;
+      chunkIndex?: number;
+    }> = [];
+
+    for (const vector of vectors) {
+      if (!vector.metadata) {
+        ctx.logger.warn('Vector missing metadata, skipping');
+        continue;
       }
-    }
-    try {
-      const vectors = await ctx.vector.search(VECTOR_STORE_NAME, dbQuery);
 
-      // Sort vectors by chunk index and concatenate text
-      const sortedVectors = vectors
-        .map(vec => {
-          const metadata = vec.metadata;
-          if (!metadata || typeof metadata.chunkIndex !== 'number' || typeof metadata.text !== 'string') {
-            ctx.logger.warn('Invalid chunk metadata structure for path %s', path);
-            return null;
-          }
-          return {
-            metadata,
-            index: metadata.chunkIndex as number
-          };
-        })
-        .filter(item => item !== null)
-        .sort((a, b) => a.index - b.index);
-
-      const fullText = sortedVectors
-        .map(vec => vec.metadata.text)
-        .join('\n\n');
-
-      return fullText;
-    } catch (err) {
-      ctx.logger.error('Error retrieving document by path %s: %o', path, err);
-      return '';
+      const path = typeof vector.metadata.path === 'string' ? vector.metadata.path : undefined;
+      const text = typeof vector.metadata.text === 'string' ? vector.metadata.text : '';
+      const chunkIndex = typeof vector.metadata.chunkIndex === 'number' ? vector.metadata.chunkIndex : undefined;
+
+      if (!path) {
+        ctx.logger.warn('Vector metadata path is not a string, skipping');
+        continue;
+      }
+
+      const relevanceScore = (vector as any).similarity;
+
+      ctx.logger.debug('Vector for path %s, chunk %d: similarity=%s, relevanceScore=%s', 
+                       path, chunkIndex, (vector as any).similarity, relevanceScore);
+
+      relevantChunks.push({
+        path,
+        content: text,
+        relevanceScore,
+        chunkIndex: chunkIndex
+      });
+    }
+
+    // Group chunks by path
+    const chunksByPath = new Map<string, Array<{
+      path: string;
+      content: string;
+      relevanceScore?: number;
+      chunkIndex?: number;
+    }>>();
+
+    for (const chunk of relevantChunks) {
+      if (!chunksByPath.has(chunk.path)) {
+        chunksByPath.set(chunk.path, []);
+      }
+      const pathChunks = chunksByPath.get(chunk.path);
+      if (pathChunks) {
+        pathChunks.push(chunk);
+      }
     }
-  }
+
+    // Expand each path group together
+    const relevantDocs: RelevantDoc[] = [];
+
+    for (const [path, pathChunks] of chunksByPath) {
+      const expandedDoc = await expandPathGroup(ctx, path, pathChunks);
+      if (expandedDoc) {
+        relevantDocs.push(expandedDoc);
+      }
+    }
+
+    ctx.logger.info('Retrieved and expanded %d relevant chunks from vector search', relevantDocs.length);
+    return relevantDocs;
+  } catch (err) {
+    ctx.logger.error('Error retrieving relevant docs: %o', err);
+    return [];
+  }
+}
+

agent-docs/src/agents/doc-qa/types.ts

@@ -2,7 +2,10 @@ import { z } from 'zod';
 
 export const RelevantDocSchema = z.object({
     path: z.string(),
-    content: z.string()
+    content: z.string(),
+    relevanceScore: z.number().optional(),
+    chunkRange: z.string().optional(),
+    chunkIndex: z.number().optional()
 });
 
 export const AnswerSchema = z.object({