📦 NEW: Dynamically set document metadata (#137)

* 📦 NEW: Dynamically set document metadata

* 👌 IMPROVE: Code

* 🐛 FIX: Reference errors

Author: msaaddev

examples/nodejs/baseai/memory/ai-agent-memory/index.ts

@@ -10,6 +10,16 @@ const memoryAiAgentMemory = (): MemoryI => ({
 		deployedAt: '',
 		embeddedAt: '',
 	},
+	documents: {
+		meta: doc => {
+			// generate a URL for each document
+			const url = `https://example.com/${doc.path}`;
+			return {
+				url,
+				name: doc.name,
+			};
+		},
+	},
 });
 
 export default memoryAiAgentMemory;

packages/baseai/src/deploy/index.ts

@@ -689,7 +689,8 @@ export async function uploadDocumentsToMemory({
 			const signedUrl = await getSignedUploadUrl({
 				documentName: doc.name,
 				memoryName: name,
-				account
+				account,
+				meta: doc.meta
 			});
 
 			const uploadResponse = await uploadDocument(signedUrl, doc.blob);
@@ -919,11 +920,13 @@ export async function listMemoryDocuments({
 async function getSignedUploadUrl({
 	documentName,
 	memoryName,
-	account
+	account,
+	meta
 }: {
 	documentName: string;
 	memoryName: string;
 	account: Account;
+	meta: Record<string, string>;
 }): Promise<string> {
 	const { uploadDocument } = getMemoryApiUrls({
 		memoryName
@@ -942,6 +945,7 @@ async function getSignedUploadUrl({
 				Authorization: `Bearer ${account.apiKey}`
 			},
 			body: JSON.stringify({
+				meta,
 				memoryName,
 				ownerLogin,
 				fileName: documentName

packages/baseai/src/utils/memory/load-memory-files.ts

@@ -5,7 +5,11 @@ import { allSupportedExtensions, MEMORYSETS } from './constants';
 import { getDocumentContent } from './get-document-content';
 import { formatDocSize } from './lib';
 import loadMemoryConfig from './load-memory-config';
-import { memoryConfigSchema, type MemoryConfigI } from 'types/memory';
+import {
+	memoryConfigSchema,
+	type DocumentConfigI,
+	type MemoryConfigI
+} from 'types/memory';
 import { execSync } from 'child_process';
 import fg from 'fast-glob';
 
@@ -14,6 +18,8 @@ export interface MemoryDocumentI {
 	size: string;
 	content: string;
 	blob: Blob;
+	path: string;
+	meta: Record<string, string>;
 }
 
 export const loadMemoryFiles = async (
@@ -24,10 +30,11 @@ export const loadMemoryFiles = async (
 
 	// useDocumentsDir
 	const useDocumentsDir = !memoryConfig || !memoryConfig.git.enabled;
+	const documentConfig = memoryConfig?.documents;
 
 	// Load files from documents directory.
 	if (useDocumentsDir) {
-		return await loadMemoryFilesFromDocsDir(memoryName);
+		return await loadMemoryFilesFromDocsDir({ memoryName, documentConfig });
 	}
 
 	// Load files from the repo.
@@ -121,12 +128,21 @@ export const loadMemoryFilesFromCustomDir = async ({
 				return null;
 			}
 
-			return {
+			const memoryFile = {
+				path: filePath,
 				name: path.basename(filePath.replace(/\//g, '-')),
 				size: formatDocSize(fileContentBlob.size),
 				content: await getDocumentContent(fileContentBlob),
 				blob: fileContentBlob
 			};
+
+			let meta = {};
+
+			if (memoryConfig?.documents?.meta) {
+				meta = memoryConfig.documents.meta(memoryFile) || {};
+			}
+
+			return { ...memoryFile, meta };
 		})
 	);
 
@@ -159,9 +175,13 @@ export const loadMemoryFilesFromCustomDir = async ({
  *    - Have unsupported file extensions.
  * 5. Returns an array of `MemoryDocumentI` objects representing the valid memory files.
  */
-export const loadMemoryFilesFromDocsDir = async (
-	memoryName: string
-): Promise<MemoryDocumentI[]> => {
+export const loadMemoryFilesFromDocsDir = async ({
+	memoryName,
+	documentConfig
+}: {
+	memoryName: string;
+	documentConfig?: DocumentConfigI;
+}): Promise<MemoryDocumentI[]> => {
 	const memoryDir = path.join(process.cwd(), 'baseai', 'memory', memoryName);
 	const memoryFilesPath = path.join(memoryDir, 'documents');
 
@@ -214,12 +234,21 @@ export const loadMemoryFilesFromDocsDir = async (
 				return null;
 			}
 
-			return {
+			const memoryFile = {
 				name: file,
+				path: filePath,
 				size: formatDocSize(fileContentBlob.size),
 				content: await getDocumentContent(fileContentBlob),
 				blob: fileContentBlob
 			};
+
+			let meta = {};
+
+			if (documentConfig?.meta) {
+				meta = documentConfig.meta(memoryFile) || {};
+			}
+
+			return { ...memoryFile, meta };
 		})
 	);
 

packages/baseai/types/memory.ts

@@ -27,14 +27,32 @@ export const gitConfigSchema = z.object({
 	embeddedAt: z.string().trim().optional().default('')
 });
 
+export const documentSchema = z.object({
+	meta: z
+		.function()
+		.args(
+			z.object({
+				name: z.string(),
+				size: z.string(),
+				content: z.string(),
+				blob: z.instanceof(Blob),
+				path: z.string(),
+			})
+		)
+		.returns(z.record(z.string()))
+		.optional()
+});
+
 export const memoryConfigSchema = z.object({
 	name: z.string(),
 	description: z.string().optional(),
-	git: gitConfigSchema
+	git: gitConfigSchema,
+	documents: documentSchema.optional()
 });
 
 export type GitConfigI = z.infer<typeof gitConfigSchema>;
 
 export type MemoryConfigI = z.infer<typeof memoryConfigSchema>;
+export type DocumentConfigI = z.infer<typeof documentSchema>;
 
 export type MemoryI = MemoryConfigI;

packages/core/types/memory.ts

@@ -6,8 +6,22 @@ export interface GitConfig {
 	embeddedAt?: string;
 }
 
+export interface MemoryDocumentI {
+	name: string;
+	size: string;
+	content: string;
+	blob: Blob;
+	path: string;
+	originalName: string;
+}
+
+export interface Document {
+	meta?: (documnet: MemoryDocumentI) => Record<string, string>;
+}
+
 export interface Memory {
 	name: string;
 	description?: string;
 	git: GitConfig;
+	documents?: Document;
 }