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-ignore - 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-ignore - 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 []
                }
            }
        })
    }
}