From ad9f99f80d76993e6d75e51a34a550c3cdff2a37 Mon Sep 17 00:00:00 2001 From: afterrburn Date: Sun, 6 Jul 2025 08:42:15 -0600 Subject: [PATCH 1/3] add totalChunks to metadata for tracing --- .../src/agents/doc-processing/chunk-mdx.ts | 2 ++ .../doc-processing/docs-orchestrator.ts | 29 ++++++++++++++----- .../agents/doc-processing/docs-processor.ts | 1 + agent-docs/src/agents/doc-processing/types.ts | 1 + 4 files changed, 26 insertions(+), 7 deletions(-) diff --git a/agent-docs/src/agents/doc-processing/chunk-mdx.ts b/agent-docs/src/agents/doc-processing/chunk-mdx.ts index 0ff24e41..5d727e05 100644 --- a/agent-docs/src/agents/doc-processing/chunk-mdx.ts +++ b/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 { return { id: crypto.randomUUID(), chunkIndex: idx, + totalChunks: chunks.length, contentType: chunk.metadata.contentType, heading: currentHeading, text: chunk.pageContent, diff --git a/agent-docs/src/agents/doc-processing/docs-orchestrator.ts b/agent-docs/src/agents/doc-processing/docs-orchestrator.ts index c78d8682..dab447b8 100644 --- a/agent-docs/src/agents/doc-processing/docs-orchestrator.ts +++ b/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) { diff --git a/agent-docs/src/agents/doc-processing/docs-processor.ts b/agent-docs/src/agents/doc-processing/docs-processor.ts index c568d26f..86ea3805 100644 --- a/agent-docs/src/agents/doc-processing/docs-processor.ts +++ b/agent-docs/src/agents/doc-processing/docs-processor.ts @@ -24,6 +24,7 @@ async function createVectorEmbedding(chunks: Chunk[]): Promise Date: Sun, 6 Jul 2025 08:42:44 -0600 Subject: [PATCH 2/3] improve RAG retrieval process --- agent-docs/config.ts | 2 +- agent-docs/src/agents/doc-qa/rag.ts | 15 +- agent-docs/src/agents/doc-qa/retriever.ts | 240 ++++++++++++++++------ agent-docs/src/agents/doc-qa/types.ts | 5 +- 4 files changed, 186 insertions(+), 76 deletions(-) diff --git a/agent-docs/config.ts b/agent-docs/config.ts index 3088c1f8..030d49d0 100644 --- a/agent-docs/config.ts +++ b/agent-docs/config.ts @@ -1,2 +1,2 @@ export const VECTOR_STORE_NAME = process.env.VECTOR_STORE_NAME || 'docs'; -export const vectorSearchNumber = 20; \ No newline at end of file +export const vectorSearchNumber = 10; \ No newline at end of file diff --git a/agent-docs/src/agents/doc-qa/rag.ts b/agent-docs/src/agents/doc-qa/rag.ts index e710e617..0aa86fad 100644 --- a/agent-docs/src/agents/doc-qa/rag.ts +++ b/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. You must first === RULES === 1. Use ONLY the content inside 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 provided to you. Keep in mind some might not be relevant, + so pick the ones that is relevant to the user's question. 4. Ambiguity handling: • When 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. - -${prompt} - - - + ${rephrasedPrompt} - + ${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, }); diff --git a/agent-docs/src/agents/doc-qa/retriever.ts b/agent-docs/src/agents/doc-qa/retriever.ts index 3c2b210d..b57c69fa 100644 --- a/agent-docs/src/agents/doc-qa/retriever.ts +++ b/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 { - 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 { + const contextWindow = 1; // Get 1 chunk before and after each chunk + const expandedChunkIndices = new Set(); + + // 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(); + // 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 { - const dbQuery = { - query: ' ', - limit: 1000, - metadata: { - path: path +} + +export async function retrieveRelevantDocs(ctx: AgentContext, prompt: string): Promise { + 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>(); + + 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); + } } - } \ No newline at end of file + + // 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 []; + } +} + diff --git a/agent-docs/src/agents/doc-qa/types.ts b/agent-docs/src/agents/doc-qa/types.ts index e1d989d4..23672e0b 100644 --- a/agent-docs/src/agents/doc-qa/types.ts +++ b/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({ From 5f8c376534dd34df6f21e24353cbb9cbfc9e71bf Mon Sep 17 00:00:00 2001 From: Seng Rith <50646727+afterrburn@users.noreply.github.com> Date: Wed, 9 Jul 2025 18:08:07 -0600 Subject: [PATCH 3/3] Apply suggestions from code review Signed-off-by: Seng Rith <50646727+afterrburn@users.noreply.github.com> --- agent-docs/src/agents/doc-qa/rag.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent-docs/src/agents/doc-qa/rag.ts b/agent-docs/src/agents/doc-qa/rag.ts index 0aa86fad..2be63873 100644 --- a/agent-docs/src/agents/doc-qa/rag.ts +++ b/agent-docs/src/agents/doc-qa/rag.ts @@ -18,7 +18,7 @@ export default async function answerQuestion(ctx: AgentContext, prompt: string) You are Agentuity's developer-documentation assistant. === CONTEXT === -Your role is to be as helpful as possible. You must first +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 tags to craft your reply. If the required information is missing, state that the docs do not cover it.