From 38ae82b370412d77d5de7980b68d15d95cca6b47 Mon Sep 17 00:00:00 2001 From: Jason Woltje Date: Fri, 13 Mar 2026 08:54:10 -0500 Subject: [PATCH] =?UTF-8?q?feat(P4-005):=20memory=20integration=20?= =?UTF-8?q?=E2=80=94=20inject=20into=20agent=20sessions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add memory tools (search, get_preferences, save_preference, save_insight) to agent sessions via Pi SDK custom tools. Agent sessions now have access to semantic memory search, preference storage, and insight capture. EmbeddingService injected into AgentService for embedding generation during tool execution. Co-Authored-By: Claude Opus 4.6 --- apps/gateway/src/agent/agent.service.ts | 12 +- apps/gateway/src/agent/tools/memory-tools.ts | 158 +++++++++++++++++++ docs/TASKS.md | 2 +- 3 files changed, 170 insertions(+), 2 deletions(-) create mode 100644 apps/gateway/src/agent/tools/memory-tools.ts diff --git a/apps/gateway/src/agent/agent.service.ts b/apps/gateway/src/agent/agent.service.ts index 9fe14f3..2e90197 100644 --- a/apps/gateway/src/agent/agent.service.ts +++ b/apps/gateway/src/agent/agent.service.ts @@ -7,11 +7,15 @@ import { type ToolDefinition, } from '@mariozechner/pi-coding-agent'; import type { Brain } from '@mosaic/brain'; +import type { Memory } from '@mosaic/memory'; import { BRAIN } from '../brain/brain.tokens.js'; +import { MEMORY } from '../memory/memory.tokens.js'; +import { EmbeddingService } from '../memory/embedding.service.js'; import { CoordService } from '../coord/coord.service.js'; import { ProviderService } from './provider.service.js'; import { createBrainTools } from './tools/brain-tools.js'; import { createCoordTools } from './tools/coord-tools.js'; +import { createMemoryTools } from './tools/memory-tools.js'; import type { SessionInfoDto } from './session.dto.js'; export interface AgentSessionOptions { @@ -42,9 +46,15 @@ export class AgentService implements OnModuleDestroy { constructor( @Inject(ProviderService) private readonly providerService: ProviderService, @Inject(BRAIN) private readonly brain: Brain, + @Inject(MEMORY) private readonly memory: Memory, + @Inject(EmbeddingService) private readonly embeddingService: EmbeddingService, @Inject(CoordService) private readonly coordService: CoordService, ) { - this.customTools = [...createBrainTools(brain), ...createCoordTools(coordService)]; + this.customTools = [ + ...createBrainTools(brain), + ...createCoordTools(coordService), + ...createMemoryTools(memory, embeddingService.available ? embeddingService : null), + ]; this.logger.log(`Registered ${this.customTools.length} custom tools`); } diff --git a/apps/gateway/src/agent/tools/memory-tools.ts b/apps/gateway/src/agent/tools/memory-tools.ts new file mode 100644 index 0000000..c8a3b43 --- /dev/null +++ b/apps/gateway/src/agent/tools/memory-tools.ts @@ -0,0 +1,158 @@ +import { Type } from '@sinclair/typebox'; +import type { ToolDefinition } from '@mariozechner/pi-coding-agent'; +import type { Memory } from '@mosaic/memory'; +import type { EmbeddingProvider } from '@mosaic/memory'; + +export function createMemoryTools( + memory: Memory, + embeddingProvider: EmbeddingProvider | null, +): ToolDefinition[] { + const searchMemory: ToolDefinition = { + name: 'memory_search', + label: 'Search Memory', + description: + 'Search across stored insights and knowledge using natural language. Returns semantically similar results.', + parameters: Type.Object({ + userId: Type.String({ description: 'User ID to search memory for' }), + query: Type.String({ description: 'Natural language search query' }), + limit: Type.Optional(Type.Number({ description: 'Max results (default 5)' })), + }), + async execute(_toolCallId, params) { + const { userId, query, limit } = params as { + userId: string; + query: string; + limit?: number; + }; + + if (!embeddingProvider) { + return { + content: [ + { + type: 'text' as const, + text: 'Semantic search unavailable — no embedding provider configured', + }, + ], + details: undefined, + }; + } + + const embedding = await embeddingProvider.embed(query); + const results = await memory.insights.searchByEmbedding(userId, embedding, limit ?? 5); + return { + content: [{ type: 'text' as const, text: JSON.stringify(results, null, 2) }], + details: undefined, + }; + }, + }; + + const getPreferences: ToolDefinition = { + name: 'memory_get_preferences', + label: 'Get User Preferences', + description: 'Retrieve stored preferences for a user.', + parameters: Type.Object({ + userId: Type.String({ description: 'User ID' }), + category: Type.Optional( + Type.String({ + description: 'Filter by category: communication, coding, workflow, appearance, general', + }), + ), + }), + async execute(_toolCallId, params) { + const { userId, category } = params as { userId: string; category?: string }; + type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general'; + const prefs = category + ? await memory.preferences.findByUserAndCategory(userId, category as Cat) + : await memory.preferences.findByUser(userId); + return { + content: [{ type: 'text' as const, text: JSON.stringify(prefs, null, 2) }], + details: undefined, + }; + }, + }; + + const savePreference: ToolDefinition = { + name: 'memory_save_preference', + label: 'Save User Preference', + description: + 'Store a learned user preference (e.g., "prefers tables over paragraphs", "timezone: America/Chicago").', + parameters: Type.Object({ + userId: Type.String({ description: 'User ID' }), + key: Type.String({ description: 'Preference key' }), + value: Type.String({ description: 'Preference value (JSON string)' }), + category: Type.Optional( + Type.String({ + description: 'Category: communication, coding, workflow, appearance, general', + }), + ), + }), + async execute(_toolCallId, params) { + const { userId, key, value, category } = params as { + userId: string; + key: string; + value: string; + category?: string; + }; + type Cat = 'communication' | 'coding' | 'workflow' | 'appearance' | 'general'; + let parsedValue: unknown; + try { + parsedValue = JSON.parse(value); + } catch { + parsedValue = value; + } + const pref = await memory.preferences.upsert({ + userId, + key, + value: parsedValue, + category: (category as Cat) ?? 'general', + source: 'agent', + }); + return { + content: [{ type: 'text' as const, text: JSON.stringify(pref, null, 2) }], + details: undefined, + }; + }, + }; + + const saveInsight: ToolDefinition = { + name: 'memory_save_insight', + label: 'Save Insight', + description: + 'Store a learned insight, decision, or knowledge extracted from the current interaction.', + parameters: Type.Object({ + userId: Type.String({ description: 'User ID' }), + content: Type.String({ description: 'The insight or knowledge to store' }), + category: Type.Optional( + Type.String({ + description: 'Category: decision, learning, preference, fact, pattern, general', + }), + ), + }), + async execute(_toolCallId, params) { + const { userId, content, category } = params as { + userId: string; + content: string; + category?: string; + }; + type Cat = 'decision' | 'learning' | 'preference' | 'fact' | 'pattern' | 'general'; + + let embedding: number[] | null = null; + if (embeddingProvider) { + embedding = await embeddingProvider.embed(content); + } + + const insight = await memory.insights.create({ + userId, + content, + embedding, + source: 'agent', + category: (category as Cat) ?? 'learning', + }); + return { + content: [{ type: 'text' as const, text: JSON.stringify(insight, null, 2) }], + details: undefined, + }; + }, + }; + + return [searchMemory, getPreferences, savePreference, saveInsight]; +} diff --git a/docs/TASKS.md b/docs/TASKS.md index 89c9844..811a813 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -41,7 +41,7 @@ | P4-002 | in-progress | Phase 4 | Semantic search — pgvector embeddings + search API | — | #35 | | P4-003 | in-progress | Phase 4 | @mosaic/log — log ingest, parsing, tiered storage | — | #36 | | P4-004 | in-progress | Phase 4 | Summarization pipeline — Haiku-tier LLM + cron | — | #37 | -| P4-005 | not-started | Phase 4 | Memory integration — inject into agent sessions | — | #38 | +| P4-005 | in-progress | Phase 4 | Memory integration — inject into agent sessions | — | #38 | | P4-006 | not-started | Phase 4 | Skill management — catalog, install, config | — | #39 | | P4-007 | not-started | Phase 4 | Verify Phase 4 — memory + log pipeline working | — | #40 | | P5-001 | not-started | Phase 5 | Plugin host — gateway plugin loading + channel interface | — | #41 |