my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import { v4 as uuidv4 } from "uuid";
#2	import { createHash } from "crypto";
#3	import {
#4	MemoryConfig,
#5	MemoryConfigSchema,
#6	MemoryItem,
#7	Message,
#8	SearchFilters,
#9	SearchResult,
#10	} from "../types";
#11	import {
#12	EmbedderFactory,
#13	LLMFactory,
#14	VectorStoreFactory,
#15	HistoryManagerFactory,
#16	} from "../utils/factory";
#17	import {
#18	FactRetrievalSchema,
#19	getFactRetrievalMessages,
#20	getUpdateMemoryMessages,
#21	parseMessages,
#22	removeCodeBlocks,
#23	} from "../prompts";
#24	import { DummyHistoryManager } from "../storage/DummyHistoryManager";
#25	import { Embedder } from "../embeddings/base";
#26	import { LLM } from "../llms/base";
#27	import { VectorStore } from "../vector_stores/base";
#28	import { ConfigManager } from "../config/manager";
#29	import { MemoryGraph } from "./graph_memory";
#30	import {
#31	AddMemoryOptions,
#32	SearchMemoryOptions,
#33	DeleteAllMemoryOptions,
#34	GetAllMemoryOptions,
#35	} from "./memory.types";
#36	import { parse_vision_messages } from "../utils/memory";
#37	import { HistoryManager } from "../storage/base";
#38	import { captureClientEvent } from "../utils/telemetry";
#39
#40	export class Memory {
#41	private config: MemoryConfig;
#42	private customPrompt: string \| undefined;
#43	private embedder: Embedder;
#44	private vectorStore: VectorStore;
#45	private llm: LLM;
#46	private db: HistoryManager;
#47	private collectionName: string \| undefined;
#48	private apiVersion: string;
#49	private graphMemory?: MemoryGraph;
#50	private enableGraph: boolean;
#51	telemetryId: string;
#52
#53	constructor(config: Partial<MemoryConfig> = {}) {
#54	// Merge and validate config
#55	this.config = ConfigManager.mergeConfig(config);
#56
#57	this.customPrompt = this.config.customPrompt;
#58	this.embedder = EmbedderFactory.create(
#59	this.config.embedder.provider,
#60	this.config.embedder.config,
#61	);
#62	this.vectorStore = VectorStoreFactory.create(
#63	this.config.vectorStore.provider,
#64	this.config.vectorStore.config,
#65	);
#66	this.llm = LLMFactory.create(
#67	this.config.llm.provider,
#68	this.config.llm.config,
#69	);
#70	if (this.config.disableHistory) {
#71	this.db = new DummyHistoryManager();
#72	} else {
#73	const defaultConfig = {
#74	provider: "sqlite",
#75	config: {
#76	historyDbPath: this.config.historyDbPath \|\| ":memory:",
#77	},
#78	};
#79
#80	this.db =
#81	this.config.historyStore && !this.config.disableHistory
#82	? HistoryManagerFactory.create(
#83	this.config.historyStore.provider,
#84	this.config.historyStore,
#85	)
#86	: HistoryManagerFactory.create("sqlite", defaultConfig);
#87	}
#88
#89	this.collectionName = this.config.vectorStore.config.collectionName;
#90	this.apiVersion = this.config.version \|\| "v1.0";
#91	this.enableGraph = this.config.enableGraph \|\| false;
#92	this.telemetryId = "anonymous";
#93
#94	// Initialize graph memory if configured
#95	if (this.enableGraph && this.config.graphStore) {
#96	this.graphMemory = new MemoryGraph(this.config);
#97	}
#98
#99	// Initialize telemetry if vector store is initialized
#100	this._initializeTelemetry();
#101	}
#102
#103	private async _initializeTelemetry() {
#104	try {
#105	await this._getTelemetryId();
#106
#107	// Capture initialization event
#108	await captureClientEvent("init", this, {
#109	api_version: this.apiVersion,
#110	client_type: "Memory",
#111	collection_name: this.collectionName,
#112	enable_graph: this.enableGraph,
#113	});
#114	} catch (error) {}
#115	}
#116
#117	private async _getTelemetryId() {
#118	try {
#119	if (
#120	!this.telemetryId \|\|
#121	this.telemetryId === "anonymous" \|\|
#122	this.telemetryId === "anonymous-supabase"
#123	) {
#124	this.telemetryId = await this.vectorStore.getUserId();
#125	}
#126	return this.telemetryId;
#127	} catch (error) {
#128	this.telemetryId = "anonymous";
#129	return this.telemetryId;
#130	}
#131	}
#132
#133	private async _captureEvent(methodName: string, additionalData = {}) {
#134	try {
#135	await this._getTelemetryId();
#136	await captureClientEvent(methodName, this, {
#137	...additionalData,
#138	api_version: this.apiVersion,
#139	collection_name: this.collectionName,
#140	});
#141	} catch (error) {
#142	console.error(`Failed to capture ${methodName} event:`, error);
#143	}
#144	}
#145
#146	static fromConfig(configDict: Record<string, any>): Memory {
#147	try {
#148	const config = MemoryConfigSchema.parse(configDict);
#149	return new Memory(config);
#150	} catch (e) {
#151	console.error("Configuration validation error:", e);
#152	throw e;
#153	}
#154	}
#155
#156	async add(
#157	messages: string \| Message[],
#158	config: AddMemoryOptions,
#159	): Promise<SearchResult> {
#160	await this._captureEvent("add", {
#161	message_count: Array.isArray(messages) ? messages.length : 1,
#162	has_metadata: !!config.metadata,
#163	has_filters: !!config.filters,
#164	infer: config.infer,
#165	});
#166	const {
#167	userId,
#168	agentId,
#169	runId,
#170	metadata = {},
#171	filters = {},
#172	infer = true,
#173	} = config;
#174
#175	if (userId) filters.userId = metadata.userId = userId;
#176	if (agentId) filters.agentId = metadata.agentId = agentId;
#177	if (runId) filters.runId = metadata.runId = runId;
#178
#179	if (!filters.userId && !filters.agentId && !filters.runId) {
#180	throw new Error(
#181	"One of the filters: userId, agentId or runId is required!",
#182	);
#183	}
#184
#185	const parsedMessages = Array.isArray(messages)
#186	? (messages as Message[])
#187	: [{ role: "user", content: messages }];
#188
#189	const final_parsedMessages = await parse_vision_messages(parsedMessages);
#190
#191	// Add to vector store
#192	const vectorStoreResult = await this.addToVectorStore(
#193	final_parsedMessages,
#194	metadata,
#195	filters,
#196	infer,
#197	);
#198
#199	// Add to graph store if available
#200	let graphResult;
#201	if (this.graphMemory) {
#202	try {
#203	graphResult = await this.graphMemory.add(
#204	final_parsedMessages.map((m) => m.content).join("\n"),
#205	filters,
#206	);
#207	} catch (error) {
#208	console.error("Error adding to graph memory:", error);
#209	}
#210	}
#211
#212	return {
#213	results: vectorStoreResult,
#214	relations: graphResult?.relations,
#215	};
#216	}
#217
#218	private async addToVectorStore(
#219	messages: Message[],
#220	metadata: Record<string, any>,
#221	filters: SearchFilters,
#222	infer: boolean,
#223	): Promise<MemoryItem[]> {
#224	if (!infer) {
#225	const returnedMemories: MemoryItem[] = [];
#226	for (const message of messages) {
#227	if (message.content === "system") {
#228	continue;
#229	}
#230	const memoryId = await this.createMemory(
#231	message.content as string,
#232	{},
#233	metadata,
#234	);
#235	returnedMemories.push({
#236	id: memoryId,
#237	memory: message.content as string,
#238	metadata: { event: "ADD" },
#239	});
#240	}
#241	return returnedMemories;
#242	}
#243	const parsedMessages = messages.map((m) => m.content).join("\n");
#244
#245	const [systemPrompt, userPrompt] = this.customPrompt
#246	? [
#247	this.customPrompt.toLowerCase().includes("json")
#248	? this.customPrompt
#249	: `${this.customPrompt}\n\nYou MUST return a valid JSON object with a 'facts' key containing an array of strings.`,
#250	`Input:\n${parsedMessages}`,
#251	]
#252	: getFactRetrievalMessages(parsedMessages);
#253
#254	const response = await this.llm.generateResponse(
#255	[
#256	{ role: "system", content: systemPrompt },
#257	{ role: "user", content: userPrompt },
#258	],
#259	{ type: "json_object" },
#260	);
#261
#262	const cleanResponse = removeCodeBlocks(response as string);
#263	let facts: string[] = [];
#264	try {
#265	const parsed = FactRetrievalSchema.parse(JSON.parse(cleanResponse));
#266	facts = parsed.facts;
#267	} catch (e) {
#268	console.error(
#269	"Failed to parse facts from LLM response:",
#270	cleanResponse,
#271	e,
#272	);
#273	facts = [];
#274	}
#275
#276	// Get embeddings for new facts
#277	const newMessageEmbeddings: Record<string, number[]> = {};
#278	const retrievedOldMemory: Array<{ id: string; text: string }> = [];
#279
#280	// Create embeddings and search for similar memories
#281	for (const fact of facts) {
#282	const embedding = await this.embedder.embed(fact);
#283	newMessageEmbeddings[fact] = embedding;
#284
#285	const existingMemories = await this.vectorStore.search(
#286	embedding,
#287	5,
#288	filters,
#289	);
#290	for (const mem of existingMemories) {
#291	retrievedOldMemory.push({ id: mem.id, text: mem.payload.data });
#292	}
#293	}
#294
#295	// Remove duplicates from old memories
#296	const uniqueOldMemories = retrievedOldMemory.filter(
#297	(mem, index) =>
#298	retrievedOldMemory.findIndex((m) => m.id === mem.id) === index,
#299	);
#300
#301	// Create UUID mapping for handling UUID hallucinations
#302	const tempUuidMapping: Record<string, string> = {};
#303	uniqueOldMemories.forEach((item, idx) => {
#304	tempUuidMapping[String(idx)] = item.id;
#305	uniqueOldMemories[idx].id = String(idx);
#306	});
#307
#308	// Get memory update decisions
#309	const updatePrompt = getUpdateMemoryMessages(uniqueOldMemories, facts);
#310
#311	const updateResponse = await this.llm.generateResponse(
#312	[{ role: "user", content: updatePrompt }],
#313	{ type: "json_object" },
#314	);
#315
#316	const cleanUpdateResponse = removeCodeBlocks(updateResponse as string);
#317	let memoryActions: any[] = [];
#318	try {
#319	memoryActions = JSON.parse(cleanUpdateResponse).memory \|\| [];
#320	} catch (e) {
#321	console.error(
#322	"Failed to parse memory actions from LLM response:",
#323	cleanUpdateResponse,
#324	e,
#325	);
#326	memoryActions = [];
#327	}
#328
#329	// Process memory actions
#330	const results: MemoryItem[] = [];
#331	for (const action of memoryActions) {
#332	try {
#333	switch (action.event) {
#334	case "ADD": {
#335	const memoryId = await this.createMemory(
#336	action.text,
#337	newMessageEmbeddings,
#338	metadata,
#339	);
#340	results.push({
#341	id: memoryId,
#342	memory: action.text,
#343	metadata: { event: action.event },
#344	});
#345	break;
#346	}
#347	case "UPDATE": {
#348	const realMemoryId = tempUuidMapping[action.id];
#349	await this.updateMemory(
#350	realMemoryId,
#351	action.text,
#352	newMessageEmbeddings,
#353	metadata,
#354	);
#355	results.push({
#356	id: realMemoryId,
#357	memory: action.text,
#358	metadata: {
#359	event: action.event,
#360	previousMemory: action.old_memory,
#361	},
#362	});
#363	break;
#364	}
#365	case "DELETE": {
#366	const realMemoryId = tempUuidMapping[action.id];
#367	await this.deleteMemory(realMemoryId);
#368	results.push({
#369	id: realMemoryId,
#370	memory: action.text,
#371	metadata: { event: action.event },
#372	});
#373	break;
#374	}
#375	}
#376	} catch (error) {
#377	console.error(`Error processing memory action: ${error}`);
#378	}
#379	}
#380
#381	return results;
#382	}
#383
#384	async get(memoryId: string): Promise<MemoryItem \| null> {
#385	const memory = await this.vectorStore.get(memoryId);
#386	if (!memory) return null;
#387
#388	const filters = {
#389	...(memory.payload.userId && { userId: memory.payload.userId }),
#390	...(memory.payload.agentId && { agentId: memory.payload.agentId }),
#391	...(memory.payload.runId && { runId: memory.payload.runId }),
#392	};
#393
#394	const memoryItem: MemoryItem = {
#395	id: memory.id,
#396	memory: memory.payload.data,
#397	hash: memory.payload.hash,
#398	createdAt: memory.payload.createdAt,
#399	updatedAt: memory.payload.updatedAt,
#400	metadata: {},
#401	};
#402
#403	// Add additional metadata
#404	const excludedKeys = new Set([
#405	"userId",
#406	"agentId",
#407	"runId",
#408	"hash",
#409	"data",
#410	"createdAt",
#411	"updatedAt",
#412	]);
#413	for (const [key, value] of Object.entries(memory.payload)) {
#414	if (!excludedKeys.has(key)) {
#415	memoryItem.metadata![key] = value;
#416	}
#417	}
#418
#419	return { ...memoryItem, ...filters };
#420	}
#421
#422	async search(
#423	query: string,
#424	config: SearchMemoryOptions,
#425	): Promise<SearchResult> {
#426	await this._captureEvent("search", {
#427	query_length: query.length,
#428	limit: config.limit,
#429	has_filters: !!config.filters,
#430	});
#431	const { userId, agentId, runId, limit = 100, filters = {} } = config;
#432
#433	if (userId) filters.userId = userId;
#434	if (agentId) filters.agentId = agentId;
#435	if (runId) filters.runId = runId;
#436
#437	if (!filters.userId && !filters.agentId && !filters.runId) {
#438	throw new Error(
#439	"One of the filters: userId, agentId or runId is required!",
#440	);
#441	}
#442
#443	// Search vector store
#444	const queryEmbedding = await this.embedder.embed(query);
#445	const memories = await this.vectorStore.search(
#446	queryEmbedding,
#447	limit,
#448	filters,
#449	);
#450
#451	// Search graph store if available
#452	let graphResults;
#453	if (this.graphMemory) {
#454	try {
#455	graphResults = await this.graphMemory.search(query, filters);
#456	} catch (error) {
#457	console.error("Error searching graph memory:", error);
#458	}
#459	}
#460
#461	const excludedKeys = new Set([
#462	"userId",
#463	"agentId",
#464	"runId",
#465	"hash",
#466	"data",
#467	"createdAt",
#468	"updatedAt",
#469	]);
#470	const results = memories.map((mem) => ({
#471	id: mem.id,
#472	memory: mem.payload.data,
#473	hash: mem.payload.hash,
#474	createdAt: mem.payload.createdAt,
#475	updatedAt: mem.payload.updatedAt,
#476	score: mem.score,
#477	metadata: Object.entries(mem.payload)
#478	.filter(([key]) => !excludedKeys.has(key))
#479	.reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}),
#480	...(mem.payload.userId && { userId: mem.payload.userId }),
#481	...(mem.payload.agentId && { agentId: mem.payload.agentId }),
#482	...(mem.payload.runId && { runId: mem.payload.runId }),
#483	}));
#484
#485	return {
#486	results,
#487	relations: graphResults,
#488	};
#489	}
#490
#491	async update(memoryId: string, data: string): Promise<{ message: string }> {
#492	await this._captureEvent("update", { memory_id: memoryId });
#493	const embedding = await this.embedder.embed(data);
#494	await this.updateMemory(memoryId, data, { [data]: embedding });
#495	return { message: "Memory updated successfully!" };
#496	}
#497
#498	async delete(memoryId: string): Promise<{ message: string }> {
#499	await this._captureEvent("delete", { memory_id: memoryId });
#500	await this.deleteMemory(memoryId);
#501	return { message: "Memory deleted successfully!" };
#502	}
#503
#504	async deleteAll(
#505	config: DeleteAllMemoryOptions,
#506	): Promise<{ message: string }> {
#507	await this._captureEvent("delete_all", {
#508	has_user_id: !!config.userId,
#509	has_agent_id: !!config.agentId,
#510	has_run_id: !!config.runId,
#511	});
#512	const { userId, agentId, runId } = config;
#513
#514	const filters: SearchFilters = {};
#515	if (userId) filters.userId = userId;
#516	if (agentId) filters.agentId = agentId;
#517	if (runId) filters.runId = runId;
#518
#519	if (!Object.keys(filters).length) {
#520	throw new Error(
#521	"At least one filter is required to delete all memories. If you want to delete all memories, use the `reset()` method.",
#522	);
#523	}
#524
#525	const [memories] = await this.vectorStore.list(filters);
#526	for (const memory of memories) {
#527	await this.deleteMemory(memory.id);
#528	}
#529
#530	return { message: "Memories deleted successfully!" };
#531	}
#532
#533	async history(memoryId: string): Promise<any[]> {
#534	return this.db.getHistory(memoryId);
#535	}
#536
#537	async reset(): Promise<void> {
#538	await this._captureEvent("reset");
#539	await this.db.reset();
#540
#541	// Check provider before attempting deleteCol
#542	if (this.config.vectorStore.provider.toLowerCase() !== "langchain") {
#543	try {
#544	await this.vectorStore.deleteCol();
#545	} catch (e) {
#546	console.error(
#547	`Failed to delete collection for provider '${this.config.vectorStore.provider}':`,
#548	e,
#549	);
#550	// Decide if you want to re-throw or just log
#551	}
#552	} else {
#553	console.warn(
#554	"Memory.reset(): Skipping vector store collection deletion as 'langchain' provider is used. Underlying Langchain vector store data is not cleared by this operation.",
#555	);
#556	}
#557
#558	if (this.graphMemory) {
#559	await this.graphMemory.deleteAll({ userId: "default" }); // Assuming this is okay, or needs similar check?
#560	}
#561
#562	// Re-initialize factories/clients based on the original config
#563	this.embedder = EmbedderFactory.create(
#564	this.config.embedder.provider,
#565	this.config.embedder.config,
#566	);
#567	// Re-create vector store instance - crucial for Langchain to reset wrapper state if needed
#568	this.vectorStore = VectorStoreFactory.create(
#569	this.config.vectorStore.provider,
#570	this.config.vectorStore.config, // This will pass the original client instance back
#571	);
#572	this.llm = LLMFactory.create(
#573	this.config.llm.provider,
#574	this.config.llm.config,
#575	);
#576	// Re-init DB if needed (though db.reset() likely handles its state)
#577	// Re-init Graph if needed
#578
#579	// Re-initialize telemetry
#580	this._initializeTelemetry();
#581	}
#582
#583	async getAll(config: GetAllMemoryOptions): Promise<SearchResult> {
#584	await this._captureEvent("get_all", {
#585	limit: config.limit,
#586	has_user_id: !!config.userId,
#587	has_agent_id: !!config.agentId,
#588	has_run_id: !!config.runId,
#589	});
#590	const { userId, agentId, runId, limit = 100 } = config;
#591
#592	const filters: SearchFilters = {};
#593	if (userId) filters.userId = userId;
#594	if (agentId) filters.agentId = agentId;
#595	if (runId) filters.runId = runId;
#596
#597	const [memories] = await this.vectorStore.list(filters, limit);
#598
#599	const excludedKeys = new Set([
#600	"userId",
#601	"agentId",
#602	"runId",
#603	"hash",
#604	"data",
#605	"createdAt",
#606	"updatedAt",
#607	]);
#608	const results = memories.map((mem) => ({
#609	id: mem.id,
#610	memory: mem.payload.data,
#611	hash: mem.payload.hash,
#612	createdAt: mem.payload.createdAt,
#613	updatedAt: mem.payload.updatedAt,
#614	metadata: Object.entries(mem.payload)
#615	.filter(([key]) => !excludedKeys.has(key))
#616	.reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}),
#617	...(mem.payload.userId && { userId: mem.payload.userId }),
#618	...(mem.payload.agentId && { agentId: mem.payload.agentId }),
#619	...(mem.payload.runId && { runId: mem.payload.runId }),
#620	}));
#621
#622	return { results };
#623	}
#624
#625	private async createMemory(
#626	data: string,
#627	existingEmbeddings: Record<string, number[]>,
#628	metadata: Record<string, any>,
#629	): Promise<string> {
#630	const memoryId = uuidv4();
#631	const embedding =
#632	existingEmbeddings[data] \|\| (await this.embedder.embed(data));
#633
#634	const memoryMetadata = {
#635	...metadata,
#636	data,
#637	hash: createHash("md5").update(data).digest("hex"),
#638	createdAt: new Date().toISOString(),
#639	};
#640
#641	await this.vectorStore.insert([embedding], [memoryId], [memoryMetadata]);
#642	await this.db.addHistory(
#643	memoryId,
#644	null,
#645	data,
#646	"ADD",
#647	memoryMetadata.createdAt,
#648	);
#649
#650	return memoryId;
#651	}
#652
#653	private async updateMemory(
#654	memoryId: string,
#655	data: string,
#656	existingEmbeddings: Record<string, number[]>,
#657	metadata: Record<string, any> = {},
#658	): Promise<string> {
#659	const existingMemory = await this.vectorStore.get(memoryId);
#660	if (!existingMemory) {
#661	throw new Error(`Memory with ID ${memoryId} not found`);
#662	}
#663
#664	const prevValue = existingMemory.payload.data;
#665	const embedding =
#666	existingEmbeddings[data] \|\| (await this.embedder.embed(data));
#667
#668	const newMetadata = {
#669	...metadata,
#670	data,
#671	hash: createHash("md5").update(data).digest("hex"),
#672	createdAt: existingMemory.payload.createdAt,
#673	updatedAt: new Date().toISOString(),
#674	...(existingMemory.payload.userId && {
#675	userId: existingMemory.payload.userId,
#676	}),
#677	...(existingMemory.payload.agentId && {
#678	agentId: existingMemory.payload.agentId,
#679	}),
#680	...(existingMemory.payload.runId && {
#681	runId: existingMemory.payload.runId,
#682	}),
#683	};
#684
#685	await this.vectorStore.update(memoryId, embedding, newMetadata);
#686	await this.db.addHistory(
#687	memoryId,
#688	prevValue,
#689	data,
#690	"UPDATE",
#691	newMetadata.createdAt,
#692	newMetadata.updatedAt,
#693	);
#694
#695	return memoryId;
#696	}
#697
#698	private async deleteMemory(memoryId: string): Promise<string> {
#699	const existingMemory = await this.vectorStore.get(memoryId);
#700	if (!existingMemory) {
#701	throw new Error(`Memory with ID ${memoryId} not found`);
#702	}
#703
#704	const prevValue = existingMemory.payload.data;
#705	await this.vectorStore.delete(memoryId);
#706	await this.db.addHistory(
#707	memoryId,
#708	prevValue,
#709	null,
#710	"DELETE",
#711	undefined,
#712	undefined,
#713	1,
#714	);
#715
#716	return memoryId;
#717	}
#718	}
#719

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public