Momento/memento-note/app/actions/ollama.ts

'use server'

interface OllamaModel {
    name: string
    modified_at: string
    size: number
    digest: string
    details: {
        format: string
        family: string
        families: string[]
        parameter_size: string
        quantization_level: string
    }
}

interface OllamaTagsResponse {
    models: OllamaModel[]
}

export async function getOllamaModels(baseUrl: string): Promise<{ success: boolean; models: string[]; error?: string }> {
    if (!baseUrl) {
        return { success: false, models: [], error: 'Base URL is required' }
    }

    // Ensure URL doesn't end with slash
    const cleanUrl = baseUrl.replace(/\/$/, '')

    // SSRF protection: block internal/private IPs
    try {
        const parsed = new URL(cleanUrl)
        const hostname = parsed.hostname.toLowerCase()
        const blockedHosts = ['localhost', '127.0.0.1', '0.0.0.0', '::1', '169.254.169.254']
        if (blockedHosts.includes(hostname)) {
            return { success: false, models: [], error: 'Private/internal URLs are not allowed' }
        }
        if (hostname.startsWith('10.') || hostname.startsWith('172.') || hostname.startsWith('192.168.') || hostname.startsWith('fc') || hostname.startsWith('fd')) {
            return { success: false, models: [], error: 'Private/internal URLs are not allowed' }
        }
        if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
            return { success: false, models: [], error: 'Only http/https protocols allowed' }
        }
    } catch {
        return { success: false, models: [], error: 'Invalid URL' }
    }

    try {
        const response = await fetch(`${cleanUrl}/api/tags`, {
            method: 'GET',
            headers: {
                'Content-Type': 'application/json',
            },
            // Set a reasonable timeout
            signal: AbortSignal.timeout(5000)
        })

        if (!response.ok) {
            throw new Error(`Ollama API returned ${response.status}: ${response.statusText}`)
        }

        const data = await response.json() as OllamaTagsResponse

        // Extract model names
        const modelNames = data.models?.map(m => m.name) || []

        return { success: true, models: modelNames }
    } catch (error: any) {
        console.error('Failed to fetch Ollama models:', error)
        return {
            success: false,
            models: [],
            error: error.message || 'Failed to connect to Ollama'
        }
    }
}