at-mintel/packages/payload-ai/src/tools/memoryDb.ts

import { tool } from "ai";
import { z } from "zod";
import { QdrantClient } from "@qdrant/js-client-rest";

// Qdrant initialization
// This requires the user to have Qdrant running and QDRANT_URL/QDRANT_API_KEY environment variables set
const qdrantClient = new QdrantClient({
  url: process.env.QDRANT_URL || "http://localhost:6333",
  apiKey: process.env.QDRANT_API_KEY,
});

const MEMORY_COLLECTION = "mintel_ai_memory";

// Ensure collection exists on load
async function initQdrant() {
  try {
    const res = await qdrantClient.getCollections();
    const exists = res.collections.find(
      (c: any) => c.name === MEMORY_COLLECTION,
    );
    if (!exists) {
      await qdrantClient.createCollection(MEMORY_COLLECTION, {
        vectors: {
          size: 1536, // typical embedding size, adjust based on the embedding model used
          distance: "Cosine",
        },
      });
      console.log(`Qdrant collection '${MEMORY_COLLECTION}' created.`);
    }
  } catch (error) {
    console.error("Failed to initialize Qdrant memory collection:", error);
  }
}

// Call init, but don't block
initQdrant();

/**
 * Returns memory tools for the AI SDK.
 * Note: A real implementation would require an embedding step before inserting into Qdrant.
 * For this implementation, we use a placeholder or assume the embeddings are handled
 * by a utility function, or we use Qdrant's FastEmbed (if running their specialized container).
 */
export const generateMemoryTools = (userId: string | number) => {
  return {
    save_memory: tool({
      description:
        "Save an important preference, fact, or instruction about the user to long-term memory. Only use this when explicitly asked or when it is clearly a long-term preference.",
      parameters: z.object({
        fact: z.string().describe("The fact or instruction to remember."),
        category: z
          .string()
          .optional()
          .describe(
            'An optional category like "preference", "rule", or "project_detail".',
          ),
      }),
      // @ts-expect-error - AI SDK strict mode bug
      execute: async ({
        fact,
        category,
      }: {
        fact: string;
        category?: string;
      }) => {
        // In a real scenario, you MUST generate embeddings for the 'fact' string here
        // using OpenAI or another embedding provider before inserting into Qdrant.
        // const embedding = await generateEmbedding(fact)

        try {
          // Mock embedding payload for demonstration
          const mockEmbedding = new Array(1536)
            .fill(0)
            .map(() => Math.random());

          await qdrantClient.upsert(MEMORY_COLLECTION, {
            wait: true,
            points: [
              {
                id: crypto.randomUUID(),
                vector: mockEmbedding,
                payload: {
                  userId: String(userId), // Partition memory by user
                  fact,
                  category,
                  createdAt: new Date().toISOString(),
                },
              },
            ],
          });
          return {
            success: true,
            message: `Successfully remembered: "${fact}"`,
          };
        } catch (error) {
          console.error("Qdrant save error:", error);
          return {
            success: false,
            error: "Failed to save to memory database.",
          };
        }
      },
    }),

    search_memory: tool({
      description:
        "Search the user's long-term memory for past factual context, preferences, or rules.",
      parameters: z.object({
        query: z.string().describe("The search string to find in memory."),
      }),
      // @ts-expect-error - AI SDK strict mode bug
      execute: async ({ query }: { query: string }) => {
        // Generate embedding for query
        const mockQueryEmbedding = new Array(1536)
          .fill(0)
          .map(() => Math.random());

        try {
          const results = await qdrantClient.search(MEMORY_COLLECTION, {
            vector: mockQueryEmbedding,
            limit: 5,
            filter: {
              must: [
                {
                  key: "userId",
                  match: { value: String(userId) },
                },
              ],
            },
          });

          return results.map((r: any) => r.payload?.fact || "");
        } catch (error) {
          console.error("Qdrant search error:", error);
          return [];
        }
      },
    }),
  };
};