stack/apps/gateway/src/agent/tools/memory-tools.ts

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];
}